Pytorch – nn.Transformer、nn.TransformerEncoderLayer、nn.TransformerEncoder、nn.TransformerDecoder、nn.TransformerDecoder参数详解

1 nn.Transformer

1.1 nn.Transformer定义

1.函数形式

torch.nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, activation='relu', custom_encoder=None, custom_decoder=None)

2.函数参数

• d_model：编码器或者解码器输入特征数量，默认值为512
• nhead：多头注意力模型中的头数，默认值为8
• num_encoder_layers：编码器中子编码器层数，默认值为6
• num_decoder_layers：解码器中子解码器层数，默认值为6
• dim_feedforward：前馈网络模型的维度，默认值为2048
• dropout：dropout值，默认为0.1
• activation：编码器或者解码器中间层的激活函数，relu或者gelu，默认为relu
• custom_encoder：自定义编码器，默认无
• custom_decoder：自定义解码器，默认无

1.2 nn.Transformer使用

1.函数形式

forward(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)

2.函数参数

• src：编码器的序列，必需参数，形状为(S,N,E)
• tgt：解码器的序列，必需参数，形状为(T,N,E)
• src_mask：src的附加掩码，可选参数，形状为(S,S)
• tgt_mask：tgt的附加掩码，可选参数，形状为(T,T)
• memory_task：编码器输出的附加掩码，可选参数，形状为(T,S)
• src_key_padding_mask：每个batch的scr keys的ByteTensor掩码，可选参数，形状为(N,S)
• tgt_key_padding_mask：每个batch的tgt keys的ByteTensor掩码，可选参数，形状为(N,T)
• memory_key_padding_mask：每个batch的memory keys的ByteTensor掩码，可选参数，形状为(N,S)

其中，S为源序列长度，T为目标序列长度，N为batch size，E为特征数

3.函数输出

输出形状为(T,N,E)的Tensor

1.3 简单使用

假设batch size为32，输入序列长度为10，输出序列长度为20，特征数量为512，采用6个多头注意力模型，编码器的子编码层数为10，则

transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
src = torch.rand((10, 32, 512))
tgt = torch.rand((20, 32, 512))
out = transformer_model(src, tgt)

2 nn.TransformerEncoderLayer

Transformer EncoderLayer 由 self-attn 和前馈网络组成。

2.1 nn.TransformerEncoderLayer定义

1.函数形式

torch.nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu')

2.函数参数

• d_model：输入特征数量，必需参数
• nhead：多头注意力模型中的头数，必需参数
• dim_feedforward：前馈网络模型的维度，默认值为2048
• dropout：dropout值，默认为0.1
• activation：编码器或者解码器中间层的激活函数，relu或者gelu，默认为relu

2.2 nn.TransformerEncoderLayer使用

1.函数形式

forward( src , src_mask=None , src_key_padding_mask=None )

2.函数参数

• src：编码器层的序列，必需参数
• src_mask：src 序列的掩码，可选参数
• src_key_padding_mask：每个batch的scr keys的ByteTensor掩码，可选参数

2.3 简单使用

encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
src = torch.rand(10, 32, 512)
out = encoder_layer(src)

3 nn.TransformerEncoder

3.1 nn.TransformerEncoder定义

1.函数形式

torch.nn.TransformerEncoder(encoder_layer, num_layers, norm=None)

nn.TransformerEncoder是堆叠num_layers个自编码器层数的模块

2.函数参数

• encoder_layer：nn.TransformerEncoderLayer的实例对象，必需参数

• num_layers：编码器中子编码器层数，必需参数

• norm：层规范化组件，可选参数

3.2 nn.TransformerEncoder使用

1.函数形式

forward(src, mask=None, src_key_padding_mask=None)

将输入依次通过编码器层。

2.函数参数

• src：编码器的输入序列，必需参数
• mask：src序列的掩码，可选参数
• src_key_padding_mask：每个batch的scr keys的ByteTensor掩码，可选参数，默认为None

2.函数输出

输出out与输入src具有相同的形状(S,N,E)。

3.3 简单使用

encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
src = torch.rand(10, 32, 512)
out = transformer_encoder(src)

4 nn.TransformerDecoderLayer

4.1 nn.TransformerDecoderLayer定义

1.函数形式

torch.nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu')

2.函数参数

• d_model：输入特征数量，必需参数

• nhead：多头注意力模型中的头数，必需参数

• dim_feedforward：前馈网络模型的维度，默认值为2048

• dropout：dropout值，默认为0.1

• activation：编码器或者解码器中间层的激活函数，relu或者gelu，默认为relu

4.2 nn.TransformerDecoderLayer使用

1.函数形式

forward(tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None)

2.函数参数

• tgt：解码器的序列，必需参数

• memory：编码器最后一层的序列，必需参数

• tgt_mask：tgt的附加掩码，可选参数

• memory_mask：编码器输出的附加掩码，可选参数

• tgt_key_padding_mask：每个batch的tgt keys的ByteTensor掩码，可选参数

• memory_key_padding_mask：每个batch的memory keys的ByteTensor掩码，可选参数

4.3 简单使用

decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
memory = torch.rand(10, 32, 512)
tgt = torch.rand(20, 32, 512)
out = decoder_layer(tgt, memory)

5 nn.TransformerDecoder

5.1 nn.TransformerDecoder定义

1.函数形式

torch.nn.TransformerDecoder(decoder_layer, num_layers, norm=None)

2.函数参数

• decoder_layer：nn.TransformerDecoderLayer的实例对象，必需参数

• num_layers：解码器中子解码器层数，必需参数

• norm：层规范化组件，可选参数

5.2 nn.TransformerDecoder使用

1.函数形式

forward(tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None)

2.函数参数

• tgt：解码器的序列，必需参数

• memory：编码器最后一层的序列，必需参数

• tgt_mask：tgt的附加掩码，可选参数

• memory_mask：编码器输出的附加掩码，可选参数

• tgt_key_padding_mask：每个batch的tgt keys的ByteTensor掩码，可选参数

• memory_key_padding_mask：每个batch的memory keys的ByteTensor掩码，可选参数

5.3 简单使用

decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
memory = torch.rand(10, 32, 512)
tgt = torch.rand(20, 32, 512)
out = transformer_decoder(tgt, memory)