STEP 13 / 18

학습률 스케줄과 학습 루프

원본: 학습 루프 · 옵티마이저

rate: 논문의 학습률 공식이다. 원본 노트북이 "This part is very important"라고 강조한 부분이다.

$$lrate = d_{model}^{-0.5} \cdot \min(step^{-0.5},\ step \cdot warmup^{-1.5})$$

처음 warmup 스텝 동안은 학습률을 직선으로 올리고, 그 뒤로는 $1/\sqrt{step}$에 비례해 내린다. 초반에 큰 학습률로 시작하면 Adam의 통계가 덜 쌓인 상태에서 크게 튀기 때문에 천천히 올린다. step 0에서는 0의 음수 제곱이 되므로 1로 바꿔 둔다. PyTorch의 LambdaLR에 넘기면 optimizer의 기본 학습률(lr=1)에 이 값이 곱해진다.

run_epoch: 데이터를 한 바퀴 돈다. 배치마다 forward → 손실 → backward → (accum_iter번에 한 번) optimizer.step() → scheduler.step() 순서다. 40배치마다 손실과 처리 속도를 출력한다. mode="eval"이면 역전파 없이 손실만 잰다.

TrainState, DummyOptimizer/DummyScheduler: 진행 상황을 세는 상자와, 평가할 때 run_epoch에 넘길 "아무것도 안 하는" 옵티마이저다.

실행 결과에서 볼 것

원본의 example_learning_schedule 차트다. 모델 크기와 warmup에 따라 꼭짓점의 위치와 높이가 달라진다. d_model이 작을수록 학습률이 높다. (원본은 2만 스텝을 모두 그리지만 페이지 크기를 줄이려고 20스텝마다 하나씩만 그렸다.)

코드 변화

+93 −0 · 파일 전체 449줄

이전 단계의 transformer.py와 비교했다. 초록 줄이 새로 생긴 줄, 빨간 줄이 사라진 줄이다.

GitHub에서 보기 ↗

실행 결과

실행한 코드 demo.py 보기

1	import altair as alt
2	import pandas as pd
3	import torch
4	from torch.optim.lr_scheduler import LambdaLR
5
6	from show import show
7	from transformer import rate
8
9	print("rate(step, model_size=512, factor=1, warmup=4000)")
10	for step in (1, 1000, 2000, 4000, 8000, 20000):
11	print(f" step {step:6d}: {rate(step, 512, 1, 4000):.2e}")
12	print("→ 4000스텝까지 직선으로 오르고, 그 뒤로 1/√step으로 내려간다")
13	print()
14
15
16	# 원본 노트북의 example_learning_schedule
17	def example_learning_schedule():
18	opts = [
19	[512, 1, 4000], # example 1
20	[512, 1, 8000], # example 2
21	[256, 1, 4000], # example 3
22	]
23
24	dummy_model = torch.nn.Linear(1, 1)
25	learning_rates = []
26
27	# we have 3 examples in opts list.
28	for idx, example in enumerate(opts):
29	# run 20000 epoch for each example
30	optimizer = torch.optim.Adam(
31	dummy_model.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9
32	)
33	lr_scheduler = LambdaLR(
34	optimizer=optimizer, lr_lambda=lambda step: rate(step, *example)
35	)
36	tmp = []
37	# take 20K dummy training steps, save the learning rate at each step
38	for step in range(20000):
39	tmp.append(optimizer.param_groups[0]["lr"])
40	optimizer.step()
41	lr_scheduler.step()
42	learning_rates.append(tmp)
43
44	learning_rates = torch.tensor(learning_rates)
45
46	# Enable altair to handle more than 5000 rows
47	alt.data_transformers.disable_max_rows()
48
49	opts_data = pd.concat(
50	[
51	pd.DataFrame(
52	{
53	"Learning Rate": learning_rates[warmup_idx, :],
54	"model_size:warmup": ["512:4000", "512:8000", "256:4000"][
55	warmup_idx
56	],
57	"step": range(20000),
58	}
59	)
60	for warmup_idx in [0, 1, 2]
61	]
62	)
63
64	return (
65	alt.Chart(opts_data[opts_data.step % 20 == 0])
66	.mark_line()
67	.properties(width=600)
68	.encode(x="step", y="Learning Rate", color="model_size\\:warmup:N")
69	.interactive()
70	)
71
72
73	show(example_learning_schedule(), "learning_rate")

rate(step, model_size=512, factor=1, warmup=4000)
  step      1: 1.75e-07
  step   1000: 1.75e-04
  step   2000: 3.49e-04
  step   4000: 6.99e-04
  step   8000: 4.94e-04
  step  20000: 3.13e-04
→ 4000스텝까지 직선으로 오르고, 그 뒤로 1/√step으로 내려간다

차트 불러오는 중…

1	1		import copy
2	2		import math
	3	+	import time
3	4
4	5		import torch
5	6		import torch.nn as nn

6	7		from torch.nn.functional import log_softmax
7	8
8	9
9	10		def attention(query, key, value, mask=None, dropout=None):
10	11		"Compute 'Scaled Dot Product Attention'"
11	12		d_k = query.size(-1)
12	13		scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
13	14		if mask is not None:
14	15		scores = scores.masked_fill(mask == 0, -1e9)
15	16		p_attn = scores.softmax(dim=-1)
16	17		if dropout is not None:
17	18		p_attn = dropout(p_attn)
18	19		return torch.matmul(p_attn, value), p_attn
19	20
20	21
21	22		def subsequent_mask(size):
22	23		"Mask out subsequent positions."
23	24		attn_shape = (1, size, size)
24	25		subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
25	26		torch.uint8
26	27		)
27	28		return subsequent_mask == 0
28	29
29	30
30	31		def clones(module, N):
31	32		"Produce N identical layers."
32	33		return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
33	34
34	35
35	36		class MultiHeadedAttention(nn.Module):
36	37		def __init__(self, h, d_model, dropout=0.1):
37	38		"Take in model size and number of heads."
38	39		super(MultiHeadedAttention, self).__init__()
39	40		assert d_model % h == 0
40	41		# We assume d_v always equals d_k
41	42		self.d_k = d_model // h
42	43		self.h = h
43	44		self.linears = clones(nn.Linear(d_model, d_model), 4)
44	45		self.attn = None
45	46		self.dropout = nn.Dropout(p=dropout)
46	47
47	48		def forward(self, query, key, value, mask=None):
48	49		"Implements Figure 2"
49	50		if mask is not None:
50	51		# Same mask applied to all h heads.
51	52		mask = mask.unsqueeze(1)
52	53		nbatches = query.size(0)
53	54
54	55		# 1) Do all the linear projections in batch from d_model => h x d_k
55	56		query, key, value = [
56	57		lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
57	58		for lin, x in zip(self.linears, (query, key, value))
58	59		]
59	60
60	61		# 2) Apply attention on all the projected vectors in batch.
61	62		x, self.attn = attention(
62	63		query, key, value, mask=mask, dropout=self.dropout
63	64		)
64	65
65	66		# 3) "Concat" using a view and apply a final linear.
66	67		x = (
67	68		x.transpose(1, 2)
68	69		.contiguous()
69	70		.view(nbatches, -1, self.h * self.d_k)
70	71		)
71	72		del query
72	73		del key
73	74		del value
74	75		return self.linears[-1](x)
75	76
76	77
77	78		class PositionwiseFeedForward(nn.Module):
78	79		"Implements FFN equation."
79	80
80	81		def __init__(self, d_model, d_ff, dropout=0.1):
81	82		super(PositionwiseFeedForward, self).__init__()
82	83		self.w_1 = nn.Linear(d_model, d_ff)
83	84		self.w_2 = nn.Linear(d_ff, d_model)
84	85		self.dropout = nn.Dropout(dropout)
85	86
86	87		def forward(self, x):
87	88		return self.w_2(self.dropout(self.w_1(x).relu()))
88	89
89	90
90	91		class Embeddings(nn.Module):
91	92		def __init__(self, d_model, vocab):
92	93		super(Embeddings, self).__init__()
93	94		self.lut = nn.Embedding(vocab, d_model)
94	95		self.d_model = d_model
95	96
96	97		def forward(self, x):
97	98		return self.lut(x) * math.sqrt(self.d_model)
98	99
99	100
100	101		class PositionalEncoding(nn.Module):
101	102		"Implement the PE function."
102	103
103	104		def __init__(self, d_model, dropout, max_len=5000):
104	105		super(PositionalEncoding, self).__init__()
105	106		self.dropout = nn.Dropout(p=dropout)
106	107
107	108		# Compute the positional encodings once in log space.
108	109		pe = torch.zeros(max_len, d_model)
109	110		position = torch.arange(0, max_len).unsqueeze(1)
110	111		div_term = torch.exp(
111	112		torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
112	113		)
113	114		pe[:, 0::2] = torch.sin(position * div_term)
114	115		pe[:, 1::2] = torch.cos(position * div_term)
115	116		pe = pe.unsqueeze(0)
116	117		self.register_buffer("pe", pe)
117	118
118	119		def forward(self, x):
119	120		x = x + self.pe[:, : x.size(1)].requires_grad_(False)
120	121		return self.dropout(x)
121	122
122	123
123	124		class LayerNorm(nn.Module):
124	125		"Construct a layernorm module (See citation for details)."
125	126
126	127		def __init__(self, features, eps=1e-6):
127	128		super(LayerNorm, self).__init__()
128	129		self.a_2 = nn.Parameter(torch.ones(features))
129	130		self.b_2 = nn.Parameter(torch.zeros(features))
130	131		self.eps = eps
131	132
132	133		def forward(self, x):
133	134		mean = x.mean(-1, keepdim=True)
134	135		std = x.std(-1, keepdim=True)
135	136		return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
136	137
137	138
138	139		class SublayerConnection(nn.Module):
139	140		"""
140	141		A residual connection followed by a layer norm.
141	142		Note for code simplicity the norm is first as opposed to last.
142	143		"""
143	144
144	145		def __init__(self, size, dropout):
145	146		super(SublayerConnection, self).__init__()
146	147		self.norm = LayerNorm(size)
147	148		self.dropout = nn.Dropout(dropout)
148	149
149	150		def forward(self, x, sublayer):
150	151		"Apply residual connection to any sublayer with the same size."
151	152		return x + self.dropout(sublayer(self.norm(x)))
152	153
153	154
154	155		class EncoderLayer(nn.Module):
155	156		"Encoder is made up of self-attn and feed forward (defined below)"
156	157
157	158		def __init__(self, size, self_attn, feed_forward, dropout):
158	159		super(EncoderLayer, self).__init__()
159	160		self.self_attn = self_attn
160	161		self.feed_forward = feed_forward
161	162		self.sublayer = clones(SublayerConnection(size, dropout), 2)
162	163		self.size = size
163	164
164	165		def forward(self, x, mask):
165	166		"Follow Figure 1 (left) for connections."
166	167		x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
167	168		return self.sublayer[1](x, self.feed_forward)
168	169
169	170
170	171		class Encoder(nn.Module):
171	172		"Core encoder is a stack of N layers"
172	173
173	174		def __init__(self, layer, N):
174	175		super(Encoder, self).__init__()
175	176		self.layers = clones(layer, N)
176	177		self.norm = LayerNorm(layer.size)
177	178
178	179		def forward(self, x, mask):
179	180		"Pass the input (and mask) through each layer in turn."
180	181		for layer in self.layers:
181	182		x = layer(x, mask)
182	183		return self.norm(x)
183	184
184	185
185	186		class DecoderLayer(nn.Module):
186	187		"Decoder is made of self-attn, src-attn, and feed forward (defined below)"
187	188
188	189		def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
189	190		super(DecoderLayer, self).__init__()
190	191		self.size = size
191	192		self.self_attn = self_attn
192	193		self.src_attn = src_attn
193	194		self.feed_forward = feed_forward
194	195		self.sublayer = clones(SublayerConnection(size, dropout), 3)
195	196
196	197		def forward(self, x, memory, src_mask, tgt_mask):
197	198		"Follow Figure 1 (right) for connections."
198	199		m = memory
199	200		x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
200	201		x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
201	202		return self.sublayer[2](x, self.feed_forward)
202	203
203	204
204	205		class Decoder(nn.Module):
205	206		"Generic N layer decoder with masking."
206	207
207	208		def __init__(self, layer, N):
208	209		super(Decoder, self).__init__()
209	210		self.layers = clones(layer, N)
210	211		self.norm = LayerNorm(layer.size)
211	212
212	213		def forward(self, x, memory, src_mask, tgt_mask):
213	214		for layer in self.layers:
214	215		x = layer(x, memory, src_mask, tgt_mask)
215	216		return self.norm(x)
216	217
217	218
218	219		class EncoderDecoder(nn.Module):
219	220		"""
220	221		A standard Encoder-Decoder architecture. Base for this and many
221	222		other models.
222	223		"""
223	224
224	225		def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
225	226		super(EncoderDecoder, self).__init__()
226	227		self.encoder = encoder
227	228		self.decoder = decoder
228	229		self.src_embed = src_embed
229	230		self.tgt_embed = tgt_embed
230	231		self.generator = generator
231	232
232	233		def forward(self, src, tgt, src_mask, tgt_mask):
233	234		"Take in and process masked src and target sequences."
234	235		return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
235	236
236	237		def encode(self, src, src_mask):
237	238		return self.encoder(self.src_embed(src), src_mask)
238	239
239	240		def decode(self, memory, src_mask, tgt, tgt_mask):
240	241		return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
241	242
242	243
243	244		class Generator(nn.Module):
244	245		"Define standard linear + softmax generation step."
245	246
246	247		def __init__(self, d_model, vocab):
247	248		super(Generator, self).__init__()
248	249		self.proj = nn.Linear(d_model, vocab)
249	250
250	251		def forward(self, x):
251	252		return log_softmax(self.proj(x), dim=-1)
252	253
253	254
254	255		def make_model(
255	256		src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
256	257		):
257	258		"Helper: Construct a model from hyperparameters."
258	259		c = copy.deepcopy
259	260		attn = MultiHeadedAttention(h, d_model)
260	261		ff = PositionwiseFeedForward(d_model, d_ff, dropout)
261	262		position = PositionalEncoding(d_model, dropout)
262	263		model = EncoderDecoder(
263	264		Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
264	265		Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
265	266		nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
266	267		nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
267	268		Generator(d_model, tgt_vocab),
268	269		)
269	270
270	271		# This was important from their code.
271	272		# Initialize parameters with Glorot / fan_avg.
272	273		for p in model.parameters():
273	274		if p.dim() > 1:
274	275		nn.init.xavier_uniform_(p)
275	276		return model
276	277
277	278
278	279		def greedy_decode(model, src, src_mask, max_len, start_symbol):
279	280		memory = model.encode(src, src_mask)
280	281		ys = torch.zeros(1, 1).fill_(start_symbol).type_as(src.data)
281	282		for i in range(max_len - 1):
282	283		out = model.decode(
283	284		memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
284	285		)
285	286		prob = model.generator(out[:, -1])
286	287		_, next_word = torch.max(prob, dim=1)
287	288		next_word = next_word.data[0]
288	289		ys = torch.cat(
289	290		[ys, torch.zeros(1, 1).type_as(src.data).fill_(next_word)], dim=1
290	291		)
291	292		return ys
292	293
293	294
294	295		class Batch:
295	296		"""Object for holding a batch of data with mask during training."""
296	297
297	298		def __init__(self, src, tgt=None, pad=2): # 2 = <blank>
298	299		self.src = src
299	300		self.src_mask = (src != pad).unsqueeze(-2)
300	301		if tgt is not None:
301	302		self.tgt = tgt[:, :-1]
302	303		self.tgt_y = tgt[:, 1:]
303	304		self.tgt_mask = self.make_std_mask(self.tgt, pad)
304	305		self.ntokens = (self.tgt_y != pad).data.sum()
305	306
306	307		@staticmethod
307	308		def make_std_mask(tgt, pad):
308	309		"Create a mask to hide padding and future words."
309	310		tgt_mask = (tgt != pad).unsqueeze(-2)
310	311		tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(
311	312		tgt_mask.data
312	313		)
313	314		return tgt_mask
314	315
315	316
316	317		class LabelSmoothing(nn.Module):
317	318		"Implement label smoothing."
318	319
319	320		def __init__(self, size, padding_idx, smoothing=0.0):
320	321		super(LabelSmoothing, self).__init__()
321	322		self.criterion = nn.KLDivLoss(reduction="sum")
322	323		self.padding_idx = padding_idx
323	324		self.confidence = 1.0 - smoothing
324	325		self.smoothing = smoothing
325	326		self.size = size
326	327		self.true_dist = None
327	328
328	329		def forward(self, x, target):
329	330		assert x.size(1) == self.size
330	331		true_dist = x.data.clone()
331	332		true_dist.fill_(self.smoothing / (self.size - 2))
332	333		true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
333	334		true_dist[:, self.padding_idx] = 0
334	335		mask = torch.nonzero(target.data == self.padding_idx)
335	336		if mask.dim() > 0:
336	337		true_dist.index_fill_(0, mask.squeeze(), 0.0)
337	338		self.true_dist = true_dist
338	339		return self.criterion(x, true_dist.clone().detach())
339	340
340	341
341	342		class SimpleLossCompute:
342	343		"A simple loss compute and train function."
343	344
344	345		def __init__(self, generator, criterion):
345	346		self.generator = generator
346	347		self.criterion = criterion
347	348
348	349		def __call__(self, x, y, norm):
349	350		x = self.generator(x)
350	351		sloss = (
351	352		self.criterion(
352	353		x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)
353	354		)
354	355		/ norm
355	356		)
356	357		return sloss.data * norm, sloss
	358	+
	359	+
	360	+	def rate(step, model_size, factor, warmup):
	361	+	"""
	362	+	we have to default the step to 1 for LambdaLR function
	363	+	to avoid zero raising to negative power.
	364	+	"""
	365	+	if step == 0:
	366	+	step = 1
	367	+	return factor * (
	368	+	model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))
	369	+	)
	370	+
	371	+
	372	+	class TrainState:
	373	+	"""Track number of steps, examples, and tokens processed"""
	374	+
	375	+	step: int = 0 # Steps in the current epoch
	376	+	accum_step: int = 0 # Number of gradient accumulation steps
	377	+	samples: int = 0 # total # of examples used
	378	+	tokens: int = 0 # total # of tokens processed
	379	+
	380	+
	381	+	def run_epoch(
	382	+	data_iter,
	383	+	model,
	384	+	loss_compute,
	385	+	optimizer,
	386	+	scheduler,
	387	+	mode="train",
	388	+	accum_iter=1,
	389	+	train_state=TrainState(),
	390	+	):
	391	+	"""Train a single epoch"""
	392	+	start = time.time()
	393	+	total_tokens = 0
	394	+	total_loss = 0
	395	+	tokens = 0
	396	+	n_accum = 0
	397	+	for i, batch in enumerate(data_iter):
	398	+	out = model.forward(
	399	+	batch.src, batch.tgt, batch.src_mask, batch.tgt_mask
	400	+	)
	401	+	loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens)
	402	+	# loss_node = loss_node / accum_iter
	403	+	if mode == "train" or mode == "train+log":
	404	+	loss_node.backward()
	405	+	train_state.step += 1
	406	+	train_state.samples += batch.src.shape[0]
	407	+	train_state.tokens += batch.ntokens
	408	+	if i % accum_iter == 0:
	409	+	optimizer.step()
	410	+	optimizer.zero_grad(set_to_none=True)
	411	+	n_accum += 1
	412	+	train_state.accum_step += 1
	413	+	scheduler.step()
	414	+
	415	+	total_loss += loss
	416	+	total_tokens += batch.ntokens
	417	+	tokens += batch.ntokens
	418	+	if i % 40 == 1 and (mode == "train" or mode == "train+log"):
	419	+	lr = optimizer.param_groups[0]["lr"]
	420	+	elapsed = time.time() - start
	421	+	print(
	422	+	(
	423	+	"Epoch Step: %6d \| Accumulation Step: %3d \| Loss: %6.2f "
	424	+	+ "\| Tokens / Sec: %7.1f \| Learning Rate: %6.1e"
	425	+	)
	426	+	% (i, n_accum, loss / batch.ntokens, tokens / elapsed, lr)
	427	+	)
	428	+	start = time.time()
	429	+	tokens = 0
	430	+	del loss
	431	+	del loss_node
	432	+	return total_loss / total_tokens, train_state
	433	+
	434	+
	435	+	class DummyOptimizer(torch.optim.Optimizer):
	436	+	def __init__(self):
	437	+	self.param_groups = [{"lr": 0}]
	438	+	None
	439	+
	440	+	def step(self):
	441	+	None
	442	+
	443	+	def zero_grad(self, set_to_none=False):
	444	+	None
	445	+
	446	+
	447	+	class DummyScheduler:
	448	+	def step(self):
	449	+	None