STEP 14 / 18

첫 학습: 입력을 그대로 따라 쓰기

원본: 첫 번째 예제 · 합성 데이터

모든 부품이 모였다. 진짜 번역 데이터로 가기 전에 가장 쉬운 과제로 전체가 맞물려 돌아가는지 확인한다. 입력 숫자열을 그대로 출력하는 복사 과제다.

data_gen: 1~10 사이 난수 10개짜리 문장을 만든다. 첫 토큰은 시작 기호 1로 고정하고 src와 tgt를 똑같이 둔다.
example_simple_model: 사전 크기 11, 2층짜리 모델을 만들어 Adam + rate(warmup 400)로 20 epoch 학습한 뒤, greedy_decode로 [0, 1, ..., 9]를 따라 쓰게 한다.

10단계에서 같은 크기의 모델이 무작위 토큰을 내놓던 것과 비교해 보자.

원본 노트북은 이 예제의 실행 줄(execute_example(example_simple_model))을 주석 처리해 두어서 출력이 없다. 여기서는 CPU로 실제로 돌린 결과를 싣는다(시드 0).

실행 결과에서 볼 것

epoch마다 한 줄씩 손실이 찍힌다. 2.9에서 시작해 0.1 안팎까지 내려간다.
학습률은 끝까지 오르기만 한다. 20 epoch × 20배치 = 400스텝이 warmup(400)과 같아서 내려가는 구간에 들어가지 못한다.
마지막 출력 [0, 2, 1, 3, 4, ...]는 앞쪽 두세 칸이 틀린다. 원본 README에 적힌 실행 결과도 [0, 2, 3, 2, 4, ...]처럼 앞쪽이 틀린다.

왜 앞쪽이 틀릴까

data_gen은 항상 1로 시작하고 1~10만 쓰는 문장을 만든다. 그런데 테스트 입력 [0, 1, ..., 9]는 0으로 시작한다. 0은 이 과제에서 패딩 id(Batch(src, tgt, 0), padding_idx=0)라 학습 중 한 번도 입력으로 나온 적이 없는 토큰이다. 출력 아래쪽 실험에서 같은 설정으로 다시 학습한 모델에 학습 때와 같은 모양의 입력([1, 2, ..., 10])을 넣으면 완벽하게 따라 쓴다. 무작위 문장 100개 중에서는 88개를 정확히 복사했다. 모델이 못 배운 게 아니라 학습 때 본 적 없는 입력을 받은 것이다.

코드 변화

+54 −0 · 파일 전체 503줄

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

GitHub에서 보기 ↗

실행 결과

실행한 코드 demo.py 보기

1	import torch
2
3	from transformer import data_gen, example_simple_model
4
5	torch.manual_seed(0)
6
7	print("복사 과제의 데이터 한 배치 (앞 3문장)")
8	b = next(data_gen(V=11, batch_size=3, nbatches=1))
9	print(" src ", b.src.tolist())
10	print(" tgt_y", b.tgt_y.tolist(), " ← src를 그대로 따라 쓰는 게 정답")
11	print()
12
13	print("원본의 example_simple_model 실행: N=2 모델, 20 epoch (CPU 약 2분)")
14	print("한 epoch에 한 줄씩, 그 epoch 두 번째 배치의 손실이 찍힌다")
15	print()
16	example_simple_model()
17	print()
18	print("입력 [0, 1, ..., 9]를 그대로 따라 썼다면 학습 성공")
19	print()
20	print("─" * 60)
21	print("왜 [0, 1, ..., 9]를 완벽하게 따라 쓰지 못할까?")
22	print("학습 데이터는 항상 1로 시작하고 1~10만 쓴다. 0은 이 과제에서 패딩 id라 한 번도 입력으로 나온 적이 없다.")
23	print("같은 설정으로 다시 학습해서 학습 때와 같은 모양의 입력을 넣어 본다.")
24	print()
25
26	import contextlib
27	import io
28
29	from torch.optim.lr_scheduler import LambdaLR
30
31	from transformer import (
32	DummyOptimizer, DummyScheduler, LabelSmoothing, SimpleLossCompute,
33	greedy_decode, make_model, rate, run_epoch,
34	)
35
36	V = 11
37	criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)
38	model = make_model(V, V, N=2)
39	optimizer = torch.optim.Adam(model.parameters(), lr=0.5, betas=(0.9, 0.98), eps=1e-9)
40	lr_scheduler = LambdaLR(optimizer=optimizer, lr_lambda=lambda step: rate(step, 512, 1.0, 400))
41	for epoch in range(20):
42	model.train()
43	with contextlib.redirect_stdout(io.StringIO()):
44	run_epoch(data_gen(V, 80, 20), model, SimpleLossCompute(model.generator, criterion),
45	optimizer, lr_scheduler, mode="train")
46	model.eval()
47
48
49	def copy(seq):
50	src = torch.LongTensor([seq])
51	out = greedy_decode(model, src, torch.ones(1, 1, len(seq)), max_len=len(seq), start_symbol=seq[0])
52	return out[0].tolist()
53
54
55	print(" 원본 테스트 [0..9] →", copy(list(range(10))))
56	print(" 학습과 같은 [1..10] →", copy(list(range(1, 11))))
57	ok = 0
58	for _ in range(100):
59	seq = torch.randint(1, V, (10,)).tolist()
60	seq[0] = 1
61	ok += copy(seq) == seq
62	print(f" 학습과 같은 분포의 무작위 문장 100개 중 완벽하게 따라 쓴 것: {ok}개")

복사 과제의 데이터 한 배치 (앞 3문장)
  src   [[1, 10, 4, 1, 4, 10, 8, 4, 8, 4], [1, 7, 7, 10, 9, 7, 7, 9, 5, 4], [1, 10, 2, 5, 5, 2, 10, 10, 10, 1]]
  tgt_y [[10, 4, 1, 4, 10, 8, 4, 8, 4], [7, 7, 10, 9, 7, 7, 9, 5, 4], [10, 2, 5, 5, 2, 10, 10, 10, 1]]  ← src를 그대로 따라 쓰는 게 정답

원본의 example_simple_model 실행: N=2 모델, 20 epoch (CPU 약 2분)
한 epoch에 한 줄씩, 그 epoch 두 번째 배치의 손실이 찍힌다

Epoch Step:      1 | Accumulation Step:   2 | Loss:   2.89 | Tokens / Sec:  2617.7 | Learning Rate: 5.5e-06
Epoch Step:      1 | Accumulation Step:   2 | Loss:   2.10 | Tokens / Sec:  2902.4 | Learning Rate: 6.1e-05
Epoch Step:      1 | Accumulation Step:   2 | Loss:   1.74 | Tokens / Sec:  2884.0 | Learning Rate: 1.2e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   1.39 | Tokens / Sec:  2839.5 | Learning Rate: 1.7e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.98 | Tokens / Sec:  2821.3 | Learning Rate: 2.3e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.61 | Tokens / Sec:  2695.7 | Learning Rate: 2.8e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.30 | Tokens / Sec:  2892.2 | Learning Rate: 3.4e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.25 | Tokens / Sec:  2861.2 | Learning Rate: 3.9e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.17 | Tokens / Sec:  2821.7 | Learning Rate: 4.5e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.15 | Tokens / Sec:  2907.0 | Learning Rate: 5.0e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.11 | Tokens / Sec:  2859.0 | Learning Rate: 5.6e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.13 | Tokens / Sec:  2853.5 | Learning Rate: 6.1e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.20 | Tokens / Sec:  2807.9 | Learning Rate: 6.7e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.17 | Tokens / Sec:  2822.6 | Learning Rate: 7.2e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.20 | Tokens / Sec:  2777.5 | Learning Rate: 7.8e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.11 | Tokens / Sec:  2836.7 | Learning Rate: 8.3e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.06 | Tokens / Sec:  2854.2 | Learning Rate: 8.9e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.20 | Tokens / Sec:  2859.5 | Learning Rate: 9.4e-04
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.08 | Tokens / Sec:  2889.9 | Learning Rate: 1.0e-03
Epoch Step:      1 | Accumulation Step:   2 | Loss:   0.14 | Tokens / Sec:  2726.9 | Learning Rate: 1.1e-03
tensor([[0, 2, 1, 3, 4, 5, 6, 7, 8, 9]])

입력 [0, 1, ..., 9]를 그대로 따라 썼다면 학습 성공

────────────────────────────────────────────────────────────
왜 [0, 1, ..., 9]를 완벽하게 따라 쓰지 못할까?
학습 데이터는 항상 1로 시작하고 1~10만 쓴다. 0은 이 과제에서 패딩 id라 한 번도 입력으로 나온 적이 없다.
같은 설정으로 다시 학습해서 학습 때와 같은 모양의 입력을 넣어 본다.

  원본 테스트  [0..9]  → [0, 3, 2, 3, 4, 5, 6, 7, 8, 9]
  학습과 같은 [1..10] → [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
  학습과 같은 분포의 무작위 문장 100개 중 완벽하게 따라 쓴 것: 88개


1	1		import copy
2	2		import math
3	3		import time
4	4
5	5		import torch
6	6		import torch.nn as nn
7	7		from torch.nn.functional import log_softmax
	8	+	from torch.optim.lr_scheduler import LambdaLR
8	9
9	10
10	11		def attention(query, key, value, mask=None, dropout=None):

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