STEP 07 / 18

인코더: 같은 층을 6번

원본: 인코더와 디코더 스택

Transformer 구조

그림 왼쪽이 인코더다. 지금까지 만든 부품을 조립한다.

EncoderLayer: self-attention과 FFN을 각각 SublayerConnection으로 감싼 두 단계다. self-attention이므로 query, key, value에 모두 같은 x를 넣는다.
Encoder: EncoderLayer를 clones로 N개 복제해 차례로 통과시키고, 마지막에 LayerNorm을 한 번 더 건다. Pre-LN 구조에서는 마지막 층의 출력이 정규화되지 않은 상태라 이 정규화가 필요하다.

EncoderLayer는 어텐션과 FFN 객체를 생성자 인자로 받는다. 어떤 어텐션을 쓸지는 조립하는 쪽(9단계의 make_model)이 정한다.

실행 결과에서 볼 것

6층 인코더의 파라미터는 약 1,890만 개다.
층마다 가중치가 다른 객체다(clones가 deepcopy하기 때문).
패딩 자리 값을 100배 크기의 난수로 바꿔도 진짜 토큰 4개의 출력은 완전히 같다(차이 0). 2단계의 마스크가 6층을 통과하는 동안 계속 지켜진다.

코드 변화

+31 −0 · 파일 전체 181줄

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

GitHub에서 보기 ↗

실행 결과

실행한 코드 demo.py 보기

1	import copy
2
3	import torch
4
5	from transformer import Encoder, EncoderLayer, MultiHeadedAttention, PositionwiseFeedForward
6
7	torch.manual_seed(0)
8
9	d_model, N = 512, 6
10	c = copy.deepcopy
11	attn = MultiHeadedAttention(8, d_model)
12	ff = PositionwiseFeedForward(d_model, 2048, 0.1)
13	encoder = Encoder(EncoderLayer(d_model, c(attn), c(ff), 0.1), N).eval()
14
15	print(f"인코더 층 {len(encoder.layers)}개, 층마다 서브층 {len(encoder.layers[0].sublayer)}개 (self-attn, FFN)")
16	n = sum(p.numel() for p in encoder.parameters())
17	print(f"인코더 파라미터 수: {n:,}")
18	print("층마다 별도 가중치?", encoder.layers[0].self_attn.linears[0].weight is not encoder.layers[1].self_attn.linears[0].weight)
19	print()
20
21	x = torch.randn(1, 6, d_model)
22	mask = torch.tensor([[[True, True, True, True, False, False]]]) # 뒤 2개는 패딩
23	out = encoder(x, mask)
24	print("입력", tuple(x.shape), "→ 출력", tuple(out.shape))
25	print()
26
27	print("패딩 자리 값을 마음대로 바꿔도 진짜 토큰의 출력은 그대로인가?")
28	x2 = x.clone()
29	x2[:, 4:] = torch.randn(1, 2, d_model) * 100
30	out2 = encoder(x2, mask)
31	print(f" 진짜 토큰 4개 출력 차이: {(out[:, :4] - out2[:, :4]).abs().max().item():.1e}")
32	print(" → 마스크 덕분에 어텐션이 패딩을 보지 않는다")

인코더 층 6개, 층마다 서브층 2개 (self-attn, FFN)
인코더 파라미터 수: 18,915,328
층마다 별도 가중치? True

입력 (1, 6, 512) → 출력 (1, 6, 512)

패딩 자리 값을 마음대로 바꿔도 진짜 토큰의 출력은 그대로인가?
  진짜 토큰 4개 출력 차이: 0.0e+00
  → 마스크 덕분에 어텐션이 패딩을 보지 않는다


1	1		import copy
2	2		import math
3	3
4	4		import torch
5	5		import torch.nn as nn
6	6
7	7
8	8		def attention(query, key, value, mask=None, dropout=None):
9	9		"Compute 'Scaled Dot Product Attention'"
10	10		d_k = query.size(-1)
11	11		scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
12	12		if mask is not None:
13	13		scores = scores.masked_fill(mask == 0, -1e9)
14	14		p_attn = scores.softmax(dim=-1)
15	15		if dropout is not None:
16	16		p_attn = dropout(p_attn)
17	17		return torch.matmul(p_attn, value), p_attn
18	18
19	19
20	20		def subsequent_mask(size):
21	21		"Mask out subsequent positions."
22	22		attn_shape = (1, size, size)
23	23		subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
24	24		torch.uint8
25	25		)
26	26		return subsequent_mask == 0
27	27
28	28
29	29		def clones(module, N):
30	30		"Produce N identical layers."
31	31		return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
32	32
33	33
34	34		class MultiHeadedAttention(nn.Module):
35	35		def __init__(self, h, d_model, dropout=0.1):
36	36		"Take in model size and number of heads."
37	37		super(MultiHeadedAttention, self).__init__()
38	38		assert d_model % h == 0
39	39		# We assume d_v always equals d_k
40	40		self.d_k = d_model // h
41	41		self.h = h
42	42		self.linears = clones(nn.Linear(d_model, d_model), 4)
43	43		self.attn = None
44	44		self.dropout = nn.Dropout(p=dropout)
45	45
46	46		def forward(self, query, key, value, mask=None):
47	47		"Implements Figure 2"
48	48		if mask is not None:
49	49		# Same mask applied to all h heads.
50	50		mask = mask.unsqueeze(1)
51	51		nbatches = query.size(0)
52	52
53	53		# 1) Do all the linear projections in batch from d_model => h x d_k
54	54		query, key, value = [
55	55		lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
56	56		for lin, x in zip(self.linears, (query, key, value))
57	57		]
58	58
59	59		# 2) Apply attention on all the projected vectors in batch.
60	60		x, self.attn = attention(
61	61		query, key, value, mask=mask, dropout=self.dropout
62	62		)
63	63
64	64		# 3) "Concat" using a view and apply a final linear.
65	65		x = (
66	66		x.transpose(1, 2)
67	67		.contiguous()
68	68		.view(nbatches, -1, self.h * self.d_k)
69	69		)
70	70		del query
71	71		del key
72	72		del value
73	73		return self.linears[-1](x)
74	74
75	75
76	76		class PositionwiseFeedForward(nn.Module):
77	77		"Implements FFN equation."
78	78
79	79		def __init__(self, d_model, d_ff, dropout=0.1):
80	80		super(PositionwiseFeedForward, self).__init__()
81	81		self.w_1 = nn.Linear(d_model, d_ff)
82	82		self.w_2 = nn.Linear(d_ff, d_model)
83	83		self.dropout = nn.Dropout(dropout)
84	84
85	85		def forward(self, x):
86	86		return self.w_2(self.dropout(self.w_1(x).relu()))
87	87
88	88
89	89		class Embeddings(nn.Module):
90	90		def __init__(self, d_model, vocab):
91	91		super(Embeddings, self).__init__()
92	92		self.lut = nn.Embedding(vocab, d_model)
93	93		self.d_model = d_model
94	94
95	95		def forward(self, x):
96	96		return self.lut(x) * math.sqrt(self.d_model)
97	97
98	98
99	99		class PositionalEncoding(nn.Module):
100	100		"Implement the PE function."
101	101
102	102		def __init__(self, d_model, dropout, max_len=5000):
103	103		super(PositionalEncoding, self).__init__()
104	104		self.dropout = nn.Dropout(p=dropout)
105	105
106	106		# Compute the positional encodings once in log space.
107	107		pe = torch.zeros(max_len, d_model)
108	108		position = torch.arange(0, max_len).unsqueeze(1)
109	109		div_term = torch.exp(
110	110		torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
111	111		)
112	112		pe[:, 0::2] = torch.sin(position * div_term)
113	113		pe[:, 1::2] = torch.cos(position * div_term)
114	114		pe = pe.unsqueeze(0)
115	115		self.register_buffer("pe", pe)
116	116
117	117		def forward(self, x):
118	118		x = x + self.pe[:, : x.size(1)].requires_grad_(False)
119	119		return self.dropout(x)
120	120
121	121
122	122		class LayerNorm(nn.Module):
123	123		"Construct a layernorm module (See citation for details)."
124	124
125	125		def __init__(self, features, eps=1e-6):
126	126		super(LayerNorm, self).__init__()
127	127		self.a_2 = nn.Parameter(torch.ones(features))
128	128		self.b_2 = nn.Parameter(torch.zeros(features))
129	129		self.eps = eps
130	130
131	131		def forward(self, x):
132	132		mean = x.mean(-1, keepdim=True)
133	133		std = x.std(-1, keepdim=True)
134	134		return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
135	135
136	136
137	137		class SublayerConnection(nn.Module):
138	138		"""
139	139		A residual connection followed by a layer norm.
140	140		Note for code simplicity the norm is first as opposed to last.
141	141		"""
142	142
143	143		def __init__(self, size, dropout):
144	144		super(SublayerConnection, self).__init__()
145	145		self.norm = LayerNorm(size)
146	146		self.dropout = nn.Dropout(dropout)
147	147
148	148		def forward(self, x, sublayer):
149	149		"Apply residual connection to any sublayer with the same size."
150	150		return x + self.dropout(sublayer(self.norm(x)))
	151	+
	152	+
	153	+	class EncoderLayer(nn.Module):
	154	+	"Encoder is made up of self-attn and feed forward (defined below)"
	155	+
	156	+	def __init__(self, size, self_attn, feed_forward, dropout):
	157	+	super(EncoderLayer, self).__init__()
	158	+	self.self_attn = self_attn
	159	+	self.feed_forward = feed_forward
	160	+	self.sublayer = clones(SublayerConnection(size, dropout), 2)
	161	+	self.size = size
	162	+
	163	+	def forward(self, x, mask):
	164	+	"Follow Figure 1 (left) for connections."
	165	+	x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
	166	+	return self.sublayer[1](x, self.feed_forward)
	167	+
	168	+
	169	+	class Encoder(nn.Module):
	170	+	"Core encoder is a stack of N layers"
	171	+
	172	+	def __init__(self, layer, N):
	173	+	super(Encoder, self).__init__()
	174	+	self.layers = clones(layer, N)
	175	+	self.norm = LayerNorm(layer.size)
	176	+
	177	+	def forward(self, x, mask):
	178	+	"Pass the input (and mask) through each layer in turn."
	179	+	for layer in self.layers:
	180	+	x = layer(x, mask)
	181	+	return self.norm(x)

1	import copy
2	import math
3
4	import torch
5	import torch.nn as nn
6
7
8	def attention(query, key, value, mask=None, dropout=None):
9	"Compute 'Scaled Dot Product Attention'"
10	d_k = query.size(-1)
11	scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
12	if mask is not None:
13	scores = scores.masked_fill(mask == 0, -1e9)
14	p_attn = scores.softmax(dim=-1)
15	if dropout is not None:
16	p_attn = dropout(p_attn)
17	return torch.matmul(p_attn, value), p_attn
18
19
20	def subsequent_mask(size):
21	"Mask out subsequent positions."
22	attn_shape = (1, size, size)
23	subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
24	torch.uint8
25	)
26	return subsequent_mask == 0
27
28
29	def clones(module, N):
30	"Produce N identical layers."
31	return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
32
33
34	class MultiHeadedAttention(nn.Module):
35	def __init__(self, h, d_model, dropout=0.1):
36	"Take in model size and number of heads."
37	super(MultiHeadedAttention, self).__init__()
38	assert d_model % h == 0
39	# We assume d_v always equals d_k
40	self.d_k = d_model // h
41	self.h = h
42	self.linears = clones(nn.Linear(d_model, d_model), 4)
43	self.attn = None
44	self.dropout = nn.Dropout(p=dropout)
45
46	def forward(self, query, key, value, mask=None):
47	"Implements Figure 2"
48	if mask is not None:
49	# Same mask applied to all h heads.
50	mask = mask.unsqueeze(1)
51	nbatches = query.size(0)
52
53	# 1) Do all the linear projections in batch from d_model => h x d_k
54	query, key, value = [
55	lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
56	for lin, x in zip(self.linears, (query, key, value))
57	]
58
59	# 2) Apply attention on all the projected vectors in batch.
60	x, self.attn = attention(
61	query, key, value, mask=mask, dropout=self.dropout
62	)
63
64	# 3) "Concat" using a view and apply a final linear.
65	x = (
66	x.transpose(1, 2)
67	.contiguous()
68	.view(nbatches, -1, self.h * self.d_k)
69	)
70	del query
71	del key
72	del value
73	return self.linears[-1](x)
74
75
76	class PositionwiseFeedForward(nn.Module):
77	"Implements FFN equation."
78
79	def __init__(self, d_model, d_ff, dropout=0.1):
80	super(PositionwiseFeedForward, self).__init__()
81	self.w_1 = nn.Linear(d_model, d_ff)
82	self.w_2 = nn.Linear(d_ff, d_model)
83	self.dropout = nn.Dropout(dropout)
84
85	def forward(self, x):
86	return self.w_2(self.dropout(self.w_1(x).relu()))
87
88
89	class Embeddings(nn.Module):
90	def __init__(self, d_model, vocab):
91	super(Embeddings, self).__init__()
92	self.lut = nn.Embedding(vocab, d_model)
93	self.d_model = d_model
94
95	def forward(self, x):
96	return self.lut(x) * math.sqrt(self.d_model)
97
98
99	class PositionalEncoding(nn.Module):
100	"Implement the PE function."
101
102	def __init__(self, d_model, dropout, max_len=5000):
103	super(PositionalEncoding, self).__init__()
104	self.dropout = nn.Dropout(p=dropout)
105
106	# Compute the positional encodings once in log space.
107	pe = torch.zeros(max_len, d_model)
108	position = torch.arange(0, max_len).unsqueeze(1)
109	div_term = torch.exp(
110	torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
111	)
112	pe[:, 0::2] = torch.sin(position * div_term)
113	pe[:, 1::2] = torch.cos(position * div_term)
114	pe = pe.unsqueeze(0)
115	self.register_buffer("pe", pe)
116
117	def forward(self, x):
118	x = x + self.pe[:, : x.size(1)].requires_grad_(False)
119	return self.dropout(x)
120
121
122	class LayerNorm(nn.Module):
123	"Construct a layernorm module (See citation for details)."
124
125	def __init__(self, features, eps=1e-6):
126	super(LayerNorm, self).__init__()
127	self.a_2 = nn.Parameter(torch.ones(features))
128	self.b_2 = nn.Parameter(torch.zeros(features))
129	self.eps = eps
130
131	def forward(self, x):
132	mean = x.mean(-1, keepdim=True)
133	std = x.std(-1, keepdim=True)
134	return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
135
136
137	class SublayerConnection(nn.Module):
138	"""
139	A residual connection followed by a layer norm.
140	Note for code simplicity the norm is first as opposed to last.
141	"""
142
143	def __init__(self, size, dropout):
144	super(SublayerConnection, self).__init__()
145	self.norm = LayerNorm(size)
146	self.dropout = nn.Dropout(dropout)
147
148	def forward(self, x, sublayer):
149	"Apply residual connection to any sublayer with the same size."
150	return x + self.dropout(sublayer(self.norm(x)))
151
152
153	class EncoderLayer(nn.Module):
154	"Encoder is made up of self-attn and feed forward (defined below)"
155
156	def __init__(self, size, self_attn, feed_forward, dropout):
157	super(EncoderLayer, self).__init__()
158	self.self_attn = self_attn
159	self.feed_forward = feed_forward
160	self.sublayer = clones(SublayerConnection(size, dropout), 2)
161	self.size = size
162
163	def forward(self, x, mask):
164	"Follow Figure 1 (left) for connections."
165	x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
166	return self.sublayer[1](x, self.feed_forward)
167
168
169	class Encoder(nn.Module):
170	"Core encoder is a stack of N layers"
171
172	def __init__(self, layer, N):
173	super(Encoder, self).__init__()
174	self.layers = clones(layer, N)
175	self.norm = LayerNorm(layer.size)
176
177	def forward(self, x, mask):
178	"Pass the input (and mask) through each layer in turn."
179	for layer in self.layers:
180	x = layer(x, mask)
181	return self.norm(x)