当前位置：首页 > news >正文

【模型学习之路】手写+分析Transformer

news 2025/1/15 13:01:57

手写+分析transformer

import torcha = torch.zeros((666,2,3))
b = torch.zeros((666,3,4))
print((a @ b).shape)   # 666, 2, 4a = torch.zeros((2,3))
b = torch.zeros((666,3,4))
print((a @ b).shape)  # 666, 2, 4a = torch.zeros((999,666,2,3))
b = torch.zeros((999,666,3,4))
print((a @ b).shape)  # 999, 666, 2, 4a = torch.zeros((999,1,2,3))
b = torch.zeros((999,666,3,4))
print((a @ b).shape)  # 999, 666, 2, 4a = torch.zeros((666,2,3))
b = torch.zeros((999,666,3,4))
print((a @ b).shape)   # 999, 666, 2, 4a = torch.zeros((2,3))
b = torch.zeros((999,666,3,4))
print((a @ b).shape)   # 999, 666, 2, 4

多头注意力机制

之前，我们有这样的公式：

就像从一个卷积核向多个卷积核过渡一样，我们可以使用多个注意力头，结合之前提到的多维张量的乘法，得到下面这些式子，其中h是头的数量：

其他的和单头大差不差。

多头注意力层的实现

实现如下图（注意：nn.Linear在应对三维输入时，只会在最后一维运用线性变换，如shape为(p,q,m)的输入经过Linear(m,n)会变成(p,q,n) ，四维及以上同理）。

下图的B显然是Batch_size。

分别指用来生成的矩阵，在这里显然三者都等于输入进来的，这里之所以分开写，是因为后面会出现三者不是同一个矩阵的情况。

其实就是Linear层，这个容易设计。

个人其实感觉多头注意力：

本质上几乎就是“单个更长的注意力头” （事实上，代码里也是这么用的）：

看了一下attention is all you need论文原文，应该是后者，当时写的时候是按前者写的，算了，差别不大。

在最后输出时，transformer还用到了resnet的思想，output加上了输入。

此外，我们在设计网络时，还要考虑一个叫掩码的东东，它的维度和attn（见下图）一样，在特定的场景（模型的具体应用场景）中使用，用来替换或“屏蔽”attn中特定的值，这里也简单实现一下。

先实现Attention函数，四个输入：Q K V mask

def attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,mask: torch.Tensor = None):k = K.size(-1)  # [B, h, m, k]scores = Q @ K.transpose(-2, -1) / np.sqrt(k)if mask is not None:scores = scores.masked_fill(mask, -1e9)attn = nn.Softmax(dim=-1)(scores)return attn @ V, attn

进一步，实现MultiHeadAttention类：

这里设置了n1和n2，不过在这一步，n1=n2=n，不过在后面会自然会出现两者不一样的情况。

class MultiHeadAttention(nn.Module):def __init__(self, n1, n2, h, k, v):# n1 表示encoder的， n2 表示decoder的super(MultiHeadAttention, self).__init__()self.n1 = n1self.n2 = n2self.h = hself.k = kself.v = vself.W_Q = nn.Linear(n2, h * k, bias=False)self.W_K = nn.Linear(n1, h * k, bias=False)self.W_V = nn.Linear(n1, h * v, bias=False)self.fc = nn.Linear(h * v, n2, bias=False)def forward(self, to_Q, to_K, to_V, mask=None):res = to_Qbatch_size = to_Q.size(0)Q = self.W_Q(to_Q).view(batch_size, -1, self.h, self.k).transpose(1, 2)K = self.W_K(to_K).view(batch_size, -1, self.h, self.k).transpose(1, 2)V = self.W_V(to_V).view(batch_size, -1, self.h, self.v).transpose(1, 2)if mask is not None:mask = mask.unsqueeze(1).repeat(1, self.h, 1)out_put, attn = attention(Q, K, V, mask)out_put = out_put.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.v)out_put = self.fc(out_put) + resout_put = nn.LayerNorm(self.n2)(out_put)return out_put, attn

我们封装一下，完成Multi-HeadAttention和Add&Norm部分。

Encoder

FeedForwardNet

首先是一个前馈层，这个很简单，一笔带过

class FeedForwardNet(nn.Module):def __init__(self, n, d_ff):super(FeedForwardNet, self).__init__()self.n = nself.fc = nn.Sequential(nn.Linear(n, d_ff, bias=False),nn.ReLU(),nn.Linear(d_ff, n, bias=False))def forward(self, x):out = self.fc(x) + xout = nn.LayerNorm(self.n)(out)return out

EncoderLayer

然后就可以将我们实现的MultiHeadAttention与前馈层组合起来，形成一个EncoderLayer。

class EncoderLayer(nn.Module):def __init__(self, n, h, k, v, d_ff):super(EncoderLayer, self).__init__()self.d_ff = d_ffself.multi_head_attention = MultiHeadAttention(n, n, h, k, v)self.feed_forward_net = FeedForwardNet(n, d_ff)def forward(self, x, mask=None):out, attn = self.multi_head_attention(x, x, x, mask)out = self.feed_forward_net(out)return out, attn

Encoder

对了，差点忘记写positional encoding的代码了。

以及embedding的代码，这里就不解释embedding是啥意思了。

这里有一次广播，因为不同句子（B个句子（语义块），每个句子m个字，每个字表示为长为n的向量）对应的position矩阵肯定是一样的。

class PositionalEncoding(nn.Module):# 这里直接搬运那条知乎作者的代码def __init__(self, n, dropout=0.1, max_len=5000):super(PositionalEncoding, self).__init__()self.dropout = nn.Dropout(p=dropout)pos_table = np.array([[pos / np.power(10000, 2 * i / n) for i in range(n)]if pos != 0 else np.zeros(n) for pos in range(max_len)])pos_table[1:, 0::2] = np.sin(pos_table[1:, 0::2])pos_table[1:, 1::2] = np.cos(pos_table[1:, 1::2])self.pos_table = torch.FloatTensor(pos_table)  # [m, n]def forward(self, x):  # x: [B, m, n]x += self.pos_table[:x.size(1), :]return self.dropout(x)

class Encoder(nn.Module):def __init__(self, vocab, n, h, k, v, d_ff, n_layers):super(Encoder, self).__init__()self.emb = nn.Embedding(vocab, n)  # [B, m] -> [B, m, vocab] -> [B, m, n]self.pos = PositionalEncoding(n)self.layers = nn.ModuleList([EncoderLayer(n, h, k, v, d_ff)for _ in range(n_layers)])def forward(self, x, mask=None):x = self.emb(x)x = self.pos(x)attn_lst = []for layer in self.layers:x, attn = layer(x, mask)attn_lst.append(attn)return x, attn_lst

Decoder

DecoderLayer

首先是单个decoder layer。

这里有两种MultiHeadAttention，一种输入全为decoder里的值，一种还要接收从endoder里出来的值。前者和encoder里面的一个意思，我们重点关注后者。

首先要知道一点，在NLP中，encoder和decoder输入的是不同的语言，也就是说，二者的m,n,vacab都是不一样的。在这里，encoder的表示为m,n。decoder的表示为M,N。这里只是为了写起来方便，代码中用m1 m2 n1 n2。

其中enc_attn里面长这样：

class DecoderLayer(nn.Module):def __init__(self, n1, n2, h, k, v, d_ff):super(DecoderLayer, self).__init__()self.enc_attn = MultiHeadAttention(n1, n2, h, k, v)self.dec_attn = MultiHeadAttention(n2, n2, h, k, v)self.ffn = FeedForwardNet(n2, d_ff)def forward(self, dec_in, enc_out, enc_mask=None, dec_mask=None):out, dec_attn = self.dec_attn(dec_in, dec_in, dec_in, dec_mask)out, enc_attn = self.enc_attn(out, enc_out, enc_out, enc_mask)out = self.ffn(out)return out, dec_attn, enc_attn

Decoder

然后做整个Decoder

class Decoder(nn.Module):def __init__(self, vocab, n, h, k, v, d_ff, n_layers):super(Decoder, self).__init__()self.emb = nn.Embedding(vocab, n)self.pos = PositionalEncoding(n)self.layers = nn.ModuleList([DecoderLayer(n, h, k, v, d_ff)for _ in range(n_layers)])def forward(self, x, enc_out, enc_mask=None, dec_mask=None):x = self.emb(x)x = self.pos(x)attn_lst = []for layer in self.layers:x, dec_attn, enc_attn = layer(x, enc_out, enc_mask, dec_mask)attn_lst.append((dec_attn, enc_attn))return x, attn_lst

组装Trasformer!

最后，组装！

class Transformer(nn.Module):def __init__(self, vocab1, vocab2, n1, n2, h, k, v, d_ff, n_layers,enc_mask=None, dec_mask=None, enc_dec_mask=None):super(Transformer, self).__init__()self.enc_mask = enc_maskself.dec_mask = dec_maskself.enc_dec_mask = enc_dec_maskself.encoder = Encoder(vocab1, n1, h, k, v, d_ff, n_layers)self.decoder = Decoder(vocab2, n1, n2, h, k, v, d_ff, n_layers)self.out = nn.Linear(n2, vocab2, bias=False)  # [B, m, n] -> [B, m, vocab]def forward(self, enc_in, dec_in):enc_out, enc_attn_lst = self.encoder(enc_in, self.enc_mask)dec_out, dec_attn_lst = self.decoder(dec_in, enc_out, self.enc_dec_mask, self.dec_mask)out = self.out(dec_out)return out, enc_attn_lst, dec_attn_lst

后话

测试一下

if __name__ == '__main__':test_enc_in = torch.randint(0, 50, (2, 10))  # [B, m1]test_dec_in = torch.randint(0, 50, (2, 12))  # [B, m2]model = Transformer(vocab1=100, vocab2=70, n1=32, n2=64, h=8, k=32, v=128, d_ff=128, n_layers=6,enc_mask=None, dec_mask=None, enc_dec_mask=None)test_out, enc_attn_lst, dec_attn_lst = model(test_enc_in, test_dec_in)print(test_out.shape)  # [B, m2, vocab2]# torch.Size([2, 12, 70])