Transformer详解(5)-编码器和解码器

1、Transformer编码器

import torch
from torch import nn
import copy
from norm import Norm
from multi_head_attention import MultiHeadAttention
from feed_forward import FeedForward
from pos_encoder import PositionalEncoderdef get_clones(module, N):"""Create N identical layers.Args:module: The module (layer) to be duplicated.N: The number of copies to create.Returns:A ModuleList containing N identical copies of the module."""return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])# transformer块
class EncoderLayer(nn.Module):def __init__(self, d_model=512, d_ff=2048, heads=8, dropout=0.1):super(EncoderLayer, self).__init__()self.norm_1 = Norm(d_model)self.norm_2 = Norm(d_model)self.attn = MultiHeadAttention(heads, d_model, dropout)self.ff = FeedForward(d_model, d_ff)self.dropout_1 = nn.Dropout(dropout)self.dropout_2 = nn.Dropout(dropout)def forward(self, x, mask):attn_ouput = self.attn(x, x, x, mask)attn_ouput = self.dropout_1(attn_ouput)x = x + attn_ouput  # 残差连接x = self.norm_1(x)  # 层归一化ff_output = self.ff(x)  # 前馈层ff_output = self.dropout_2(ff_output)x = x + ff_output  # 残差连接x = self.norm_2(x)  # 层归一化return xclass TransformerEncoder(nn.Module):def __init__(self, vocab_size=1000, max_seq_len=50, d_model=512, d_ff=2048, N=6, heads=8, dropout=0.1):super(TransformerEncoder, self).__init__()'''vocab_size  词典大小max_seq_len  序列最大长度d_model  词嵌入大小d_ff  前馈层隐层维度N  编码器中transformer的个数heads  多头个数dropout  dropout比例'''self.N = Nself.embed = nn.Embedding(vocab_size, d_model)self.pe = PositionalEncoder(max_seq_len, d_model)self.layers = get_clones(EncoderLayer(d_model, d_ff, heads, dropout), N)self.norm = Norm(d_model)def forward(self, src, mask=None):x = self.embed(src)  # embeddingx = self.pe(x)  # 位置编码for i in range(self.N):x = self.layers[i](x, mask)output = self.norm(x)return outputif __name__ == '__main__':# Parameterslength = 50low = 0high = 1001  # The upper bound is exclusive in torch.randint# Generate random integersrandom_tensor = torch.randint(low=low, high=high, size=(length,))vocab_size = 1000max_seq_len = 50d_model = 512d_ff = 2048heads = 8N = 2dropout = 0.1trans_encoder = TransformerEncoder(vocab_size, max_seq_len, d_model, d_ff, N, heads, dropout)output = trans_encoder(random_tensor)print(output.shape)  # torch.Size([1, 50, 512])