STEP 09 / 18

하나로 조립: EncoderDecoder와 make_model

원본: 모델 구조 · 전체 모델

원본 노트북이 맨 처음에 보여 주는 코드가 이 단계의 EncoderDecoder다. 원본을 위에서부터 읽으면 아직 아무 부품도 모르는 상태에서 이 클래스를 만나서 어디부터 봐야 할지 막막하다. 여기서는 부품을 다 만든 뒤에 만나므로 각 줄이 무엇을 부르는지 이미 알고 있다.

EncoderDecoder: 부품 다섯 개(encoder, decoder, src_embed, tgt_embed, generator)를 들고 있는 껍데기다. encode로 원문을 memory로 만들고, decode로 memory와 지금까지의 번역문을 디코더에 넣는다.
Generator: 디코더 출력(512차원)을 사전 크기로 투영하고 log_softmax를 걸어 "다음 토큰이 무엇일지"의 로그 확률을 낸다.
make_model: 하이퍼파라미터(N=6, d_model=512, d_ff=2048, h=8, dropout=0.1)로 모델을 조립한다. 어텐션·FFN·PE를 하나씩 만들어 두고 copy.deepcopy로 필요한 곳마다 복제한다. 마지막에 2차원 이상인 파라미터를 전부 Xavier(Glorot) 균등분포로 초기화한다. 원본 주석대로 "This was important"다.

실행 결과에서 볼 것

15단계에서 만들 사전 크기(독일어 8,316, 영어 6,384)로 조립하면 파라미터가 약 5,500만 개다. 디코더가 가장 크다.
encode → decode → generator를 지나며 모양이 (1, 9) → (1, 9, 512) → (1, 5, 512) → (1, 5, 6384)로 바뀐다.
초기화된 가중치의 최댓값이 Xavier 이론 한계 $\sqrt{6/(512+2048)}$와 같다.

코드 변화

+61 −0 · 파일 전체 275줄

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

GitHub에서 보기 ↗

실행 결과

실행한 코드 demo.py 보기

1	import torch
2
3	from transformer import make_model, subsequent_mask
4
5	torch.manual_seed(0)
6
7	model = make_model(src_vocab=8316, tgt_vocab=6384) # 15단계에서 만들 독일어·영어 사전 크기
8	print("EncoderDecoder의 다섯 부품")
9	for name, m in model.named_children():
10	print(f" {name:10s} {sum(p.numel() for p in m.parameters()):>12,} {m.__class__.__name__}")
11	print(f" {'합계':10s} {sum(p.numel() for p in model.parameters()):>12,}")
12	print()
13
14	model.eval()
15	src = torch.randint(4, 8316, (1, 9))
16	tgt = torch.randint(4, 6384, (1, 5))
17	src_mask = torch.ones(1, 1, 9, dtype=torch.bool)
18	memory = model.encode(src, src_mask)
19	out = model.decode(memory, src_mask, tgt, subsequent_mask(5))
20	logp = model.generator(out)
21	print("src", tuple(src.shape), "→ encode →", tuple(memory.shape))
22	print("tgt", tuple(tgt.shape), "→ decode →", tuple(out.shape))
23	print("generator →", tuple(logp.shape), " (다음 토큰의 로그 확률)")
24	print("확률 합:", [f"{v:.3f}" for v in logp.exp().sum(-1)[0].tolist()])
25	print()
26	w = model.encoder.layers[0].feed_forward.w_1.weight
27	print(f"xavier_uniform_ 초기화: w_1 값 범위 ±{w.abs().max():.4f} (이론 한계 √(6/(512+2048)) = {(6 / (512 + 2048)) ** 0.5:.4f})")

EncoderDecoder의 다섯 부품
  encoder      18,915,328  Encoder
  decoder      25,225,216  Decoder
  src_embed     4,257,792  Sequential
  tgt_embed     3,268,608  Sequential
  generator     3,274,992  Generator
  합계           54,941,936

src (1, 9) → encode → (1, 9, 512)
tgt (1, 5) → decode → (1, 5, 512)
generator → (1, 5, 6384)  (다음 토큰의 로그 확률)
확률 합: ['1.000', '1.000', '1.000', '1.000', '1.000']

xavier_uniform_ 초기화: w_1 값 범위 ±0.0484  (이론 한계 √(6/(512+2048)) = 0.0484)

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

9	10		"Compute 'Scaled Dot Product Attention'"
10	11		d_k = query.size(-1)
11	12		scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
12	13		if mask is not None:
13	14		scores = scores.masked_fill(mask == 0, -1e9)
14	15		p_attn = scores.softmax(dim=-1)
15	16		if dropout is not None:
16	17		p_attn = dropout(p_attn)
17	18		return torch.matmul(p_attn, value), p_attn
18	19
19	20
20	21		def subsequent_mask(size):
21	22		"Mask out subsequent positions."
22	23		attn_shape = (1, size, size)
23	24		subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
24	25		torch.uint8
25	26		)
26	27		return subsequent_mask == 0
27	28
28	29
29	30		def clones(module, N):
30	31		"Produce N identical layers."
31	32		return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
32	33
33	34
34	35		class MultiHeadedAttention(nn.Module):
35	36		def __init__(self, h, d_model, dropout=0.1):
36	37		"Take in model size and number of heads."
37	38		super(MultiHeadedAttention, self).__init__()
38	39		assert d_model % h == 0
39	40		# We assume d_v always equals d_k
40	41		self.d_k = d_model // h
41	42		self.h = h
42	43		self.linears = clones(nn.Linear(d_model, d_model), 4)
43	44		self.attn = None
44	45		self.dropout = nn.Dropout(p=dropout)
45	46
46	47		def forward(self, query, key, value, mask=None):
47	48		"Implements Figure 2"
48	49		if mask is not None:
49	50		# Same mask applied to all h heads.
50	51		mask = mask.unsqueeze(1)
51	52		nbatches = query.size(0)
52	53
53	54		# 1) Do all the linear projections in batch from d_model => h x d_k
54	55		query, key, value = [
55	56		lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
56	57		for lin, x in zip(self.linears, (query, key, value))
57	58		]
58	59
59	60		# 2) Apply attention on all the projected vectors in batch.
60	61		x, self.attn = attention(
61	62		query, key, value, mask=mask, dropout=self.dropout
62	63		)
63	64
64	65		# 3) "Concat" using a view and apply a final linear.
65	66		x = (
66	67		x.transpose(1, 2)
67	68		.contiguous()
68	69		.view(nbatches, -1, self.h * self.d_k)
69	70		)
70	71		del query
71	72		del key
72	73		del value
73	74		return self.linears[-1](x)
74	75
75	76
76	77		class PositionwiseFeedForward(nn.Module):
77	78		"Implements FFN equation."
78	79
79	80		def __init__(self, d_model, d_ff, dropout=0.1):
80	81		super(PositionwiseFeedForward, self).__init__()
81	82		self.w_1 = nn.Linear(d_model, d_ff)
82	83		self.w_2 = nn.Linear(d_ff, d_model)
83	84		self.dropout = nn.Dropout(dropout)
84	85
85	86		def forward(self, x):
86	87		return self.w_2(self.dropout(self.w_1(x).relu()))
87	88
88	89
89	90		class Embeddings(nn.Module):
90	91		def __init__(self, d_model, vocab):
91	92		super(Embeddings, self).__init__()
92	93		self.lut = nn.Embedding(vocab, d_model)
93	94		self.d_model = d_model
94	95
95	96		def forward(self, x):
96	97		return self.lut(x) * math.sqrt(self.d_model)
97	98
98	99
99	100		class PositionalEncoding(nn.Module):
100	101		"Implement the PE function."
101	102
102	103		def __init__(self, d_model, dropout, max_len=5000):
103	104		super(PositionalEncoding, self).__init__()
104	105		self.dropout = nn.Dropout(p=dropout)
105	106
106	107		# Compute the positional encodings once in log space.
107	108		pe = torch.zeros(max_len, d_model)
108	109		position = torch.arange(0, max_len).unsqueeze(1)
109	110		div_term = torch.exp(
110	111		torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
111	112		)
112	113		pe[:, 0::2] = torch.sin(position * div_term)
113	114		pe[:, 1::2] = torch.cos(position * div_term)
114	115		pe = pe.unsqueeze(0)
115	116		self.register_buffer("pe", pe)
116	117
117	118		def forward(self, x):
118	119		x = x + self.pe[:, : x.size(1)].requires_grad_(False)
119	120		return self.dropout(x)
120	121
121	122
122	123		class LayerNorm(nn.Module):
123	124		"Construct a layernorm module (See citation for details)."
124	125
125	126		def __init__(self, features, eps=1e-6):
126	127		super(LayerNorm, self).__init__()
127	128		self.a_2 = nn.Parameter(torch.ones(features))
128	129		self.b_2 = nn.Parameter(torch.zeros(features))
129	130		self.eps = eps
130	131
131	132		def forward(self, x):
132	133		mean = x.mean(-1, keepdim=True)
133	134		std = x.std(-1, keepdim=True)
134	135		return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
135	136
136	137
137	138		class SublayerConnection(nn.Module):
138	139		"""
139	140		A residual connection followed by a layer norm.
140	141		Note for code simplicity the norm is first as opposed to last.
141	142		"""
142	143
143	144		def __init__(self, size, dropout):
144	145		super(SublayerConnection, self).__init__()
145	146		self.norm = LayerNorm(size)
146	147		self.dropout = nn.Dropout(dropout)
147	148
148	149		def forward(self, x, sublayer):
149	150		"Apply residual connection to any sublayer with the same size."
150	151		return x + self.dropout(sublayer(self.norm(x)))
151	152
152	153
153	154		class EncoderLayer(nn.Module):
154	155		"Encoder is made up of self-attn and feed forward (defined below)"
155	156
156	157		def __init__(self, size, self_attn, feed_forward, dropout):
157	158		super(EncoderLayer, self).__init__()
158	159		self.self_attn = self_attn
159	160		self.feed_forward = feed_forward
160	161		self.sublayer = clones(SublayerConnection(size, dropout), 2)
161	162		self.size = size
162	163
163	164		def forward(self, x, mask):
164	165		"Follow Figure 1 (left) for connections."
165	166		x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
166	167		return self.sublayer[1](x, self.feed_forward)
167	168
168	169
169	170		class Encoder(nn.Module):
170	171		"Core encoder is a stack of N layers"
171	172
172	173		def __init__(self, layer, N):
173	174		super(Encoder, self).__init__()
174	175		self.layers = clones(layer, N)
175	176		self.norm = LayerNorm(layer.size)
176	177
177	178		def forward(self, x, mask):
178	179		"Pass the input (and mask) through each layer in turn."
179	180		for layer in self.layers:
180	181		x = layer(x, mask)
181	182		return self.norm(x)
182	183
183	184
184	185		class DecoderLayer(nn.Module):
185	186		"Decoder is made of self-attn, src-attn, and feed forward (defined below)"
186	187
187	188		def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
188	189		super(DecoderLayer, self).__init__()
189	190		self.size = size
190	191		self.self_attn = self_attn
191	192		self.src_attn = src_attn
192	193		self.feed_forward = feed_forward
193	194		self.sublayer = clones(SublayerConnection(size, dropout), 3)
194	195
195	196		def forward(self, x, memory, src_mask, tgt_mask):
196	197		"Follow Figure 1 (right) for connections."
197	198		m = memory
198	199		x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
199	200		x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
200	201		return self.sublayer[2](x, self.feed_forward)
201	202
202	203
203	204		class Decoder(nn.Module):
204	205		"Generic N layer decoder with masking."
205	206
206	207		def __init__(self, layer, N):
207	208		super(Decoder, self).__init__()
208	209		self.layers = clones(layer, N)
209	210		self.norm = LayerNorm(layer.size)
210	211
211	212		def forward(self, x, memory, src_mask, tgt_mask):
212	213		for layer in self.layers:
213	214		x = layer(x, memory, src_mask, tgt_mask)
214	215		return self.norm(x)
	216	+
	217	+
	218	+	class EncoderDecoder(nn.Module):
	219	+	"""
	220	+	A standard Encoder-Decoder architecture. Base for this and many
	221	+	other models.
	222	+	"""
	223	+
	224	+	def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
	225	+	super(EncoderDecoder, self).__init__()
	226	+	self.encoder = encoder
	227	+	self.decoder = decoder
	228	+	self.src_embed = src_embed
	229	+	self.tgt_embed = tgt_embed
	230	+	self.generator = generator
	231	+
	232	+	def forward(self, src, tgt, src_mask, tgt_mask):
	233	+	"Take in and process masked src and target sequences."
	234	+	return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
	235	+
	236	+	def encode(self, src, src_mask):
	237	+	return self.encoder(self.src_embed(src), src_mask)
	238	+
	239	+	def decode(self, memory, src_mask, tgt, tgt_mask):
	240	+	return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
	241	+
	242	+
	243	+	class Generator(nn.Module):
	244	+	"Define standard linear + softmax generation step."
	245	+
	246	+	def __init__(self, d_model, vocab):
	247	+	super(Generator, self).__init__()
	248	+	self.proj = nn.Linear(d_model, vocab)
	249	+
	250	+	def forward(self, x):
	251	+	return log_softmax(self.proj(x), dim=-1)
	252	+
	253	+
	254	+	def make_model(
	255	+	src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
	256	+	):
	257	+	"Helper: Construct a model from hyperparameters."
	258	+	c = copy.deepcopy
	259	+	attn = MultiHeadedAttention(h, d_model)
	260	+	ff = PositionwiseFeedForward(d_model, d_ff, dropout)
	261	+	position = PositionalEncoding(d_model, dropout)
	262	+	model = EncoderDecoder(
	263	+	Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
	264	+	Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
	265	+	nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
	266	+	nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
	267	+	Generator(d_model, tgt_vocab),
	268	+	)
	269	+
	270	+	# This was important from their code.
	271	+	# Initialize parameters with Glorot / fan_avg.
	272	+	for p in model.parameters():
	273	+	if p.dim() > 1:
	274	+	nn.init.xavier_uniform_(p)
	275	+	return model