PyTorch中Transformer模型的搭建
PyTorch中Transformer模型的搭建
PyTorch最近版本更新很快，1.2/1.3/1.4几乎是连着出，其中： 1.3/1.4版本主要是新增并完善了PyTorchMobile移动端部署模块和模型量化模块。 而1.2版中一个重要的更新就是把加入了NLP领域中炙手可热的Transformer模型，这里记录一下PyTorch中Transformer模型的用法（代码写于1.2版本，没有在1.3/1.4版本测试）。
1. 简介
也许是为了更方便地搭建Bert，GPT-2之类的NLP模型，PyTorch将Transformer相关的模型分为nn.TransformerEncoderLayer、nn.TransformerDecoderLayer、nn.LayerNorm等几个部分。搭建模型的时候不一定都会用到， 比如fastai中的Transformer模型就只用到了encoder部分，没有用到decoder。
至于WordEmbedding和PositionEncoding两个部分需要自己另外实现。
WordEmbedding可以直接使用PyTorch自带的nn.Embedding层。
PositionEncoding层的花样就多了，不同的模型下面有不同的PositionEncoding，比如Transformer的原始论文Attention is all you need中使用的是无参数的PositionEncoding， Bert中使用的是带有学习参数的PositionEncoding。
在本文中介绍的是参考Transformer原始论文实现的Sequence2sequence形式的Transformer模型。
2. Sequence2sequence形式的Transformer模型搭建：
2.1 无可学习参数的PositionEncoding层
无参数的PositionEncoding计算速度快，还可以减小整个模型的尺寸，据说在有些任务中，效果与有参数的接近。
class PositionalEncoding(nn.Module):
def __init__(self, d_model,dropout=0.1, max_len=5000):
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0).transpose(0, 1)
self.register_buffer('pe', pe)
def forward(self, x):
x = x + self.pe[:x.size(0), :]
return self.dropout(x)2.2 有可学习参数的PositionEncoding层
我曾在一个序列预测任务（非NLP）里面对比过两种PositionEncoding层，发现带有参数的PositionEncoding层效果明显比没有参数的PositionEncoding要好。
带参数的PositionEncoding层的定义更为简单，直接继承一个nn.Embedding，再续上一个dropout就可以了。因为nn.Embedding中包含了一个可以按索引取向量的权重矩阵weight。
class LearnedPositionEncoding(nn.Embedding):
def __init__(self,d_model, dropout = 0.1,max_len = 5000):
super().__init__(max_len, d_model)
self.dropout = nn.Dropout(p = dropout)
def forward(self, x):
weight = self.weight.data.unsqueeze(1)
x = x + weight[:x.size(0),:]
return self.dropout(x)2.3 Sequence2sequence模型
将embedding、position_encoding、encoder和decoder拼接起来，就可以构成一个完整的sequence2sequence形式的Transformer模型了。
class S2sTransformer(nn.Module):
def __init__(self,vocab_size,position_enc,d_model = 512,nhead = 8,num_encoder_layers=6,
num_decoder_layers=6,dim_feedforward=2048,dropout=0.1):
super(S2sTransformer,self).__init__()
# Preprocess
self.embedding = nn.Embedding(vocab_size,d_model)
self.pos_encoder_src = position_enc(d_model=512)
# tgt
self.pos_encoder_tgt = position_enc(d_model=512)
# Encoder
encoder_layer = nn.TransformerEncoderLayer(d_model,nhead,dim_feedforward,dropout)
encoder_norm = nn.LayerNorm(d_model)
self.encoder = nn.TransformerEncoder(encoder_layer,num_encoder_layers,encoder_norm)
# Decoder
decoder_layer = nn.TransformerDecoderLayer(d_model,nhead,dim_feedforward,dropout)
decoder_norm = nn.LayerNorm(d_model)
self.decoder = nn.TransformerDecoder(decoder_layer,num_decoder_layers,decoder_norm)
self.output_layer = nn.Linear(d_model,vocab_size)
self._reset_parameters()
self.d_model = d_model
self.nhead = nhead
def forward(self, src,tgt,src_mask = None,tgt_mask = None,
memory_mask = None,src_key_padding_mask = None,
tgt_key_padding_mask = None,memory_key_padding_mask = None):
# word embedding
src = self.embedding(src)
tgt = self.embedding(tgt)
# shape check
if src.size(1) != tgt.size(1):
raise RuntimeError("the batch number of src and tgt must be equal")
if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
raise RuntimeError("the feature number of src and tgt must be equal to d_model")
# position encoding
src = self.pos_encoder_src(src)
tgt = self.pos_encoder_tgt(tgt)
memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
tgt_key_padding_mask=tgt_key_padding_mask,
memory_key_padding_mask=memory_key_padding_mask)
output = self.output_layer(output)
# return output
return softmax(output,dim = 2)
def generate_square_subsequent_mask(self, sz):
r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
Unmasked positions are filled with float(0.0).
"""
mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
return mask
def _reset_parameters(self):
r"""Initiate parameters in the transformer model."""
for p in self.parameters():
if p.dim() > 1:
xavier_uniform_(p)模型搭建好了之后，就可以按照Sequence2sequence的训练方式进行训练了， 唯一需要注意的就是Transformer的forward过程是并行的，与基于RNN的Sequence2sequence模型稍有不同。
训练过程可以参考PyTorch官网提供的chatbot的教程