nn.TransformerEncoder的输出为NaN值的原因及解决方法

jq_982023-11-30 0:03

问题描述：

当使用nn.TransformerEncoder时，即如下情况：

复制代码

实例化一个TransformerEncoder 
 self.encoder_layer = nn.TransformerEncoderLayer(d_model=encoder_in_dim, nhead=encoder_head,
                                                           dim_feedforward=encoder_ffnn_dim,
                                                           batch_first=batch_first)
 self.model = nn.TransformerEncoder(self.pre_encoder_layer, num_layers=pre_encoder_layer_num)
调用：
transformer_features =  self.model(embeddings, src_key_padding_mask=src_padding_mask)

transformer_features的值为NaN

原因在于src_padding_mask的传入出现均为0/False的情况！即attention---mask出现了全1/True行

由于我们在使用MultiheadAttention做self-attention时因为batch内序列长度不一致，难免需要使用mask

以pytorch自带的torch.nn.TransformerEncoder方法为例，其forward函数如下

复制代码

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

这里的mask会送到torch.nn.TransformerEncoderLayer的forward函数：

python 复制代码

def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        src2 = self.self_attn(src, src, src, attn_mask=src_mask,key_padding_mask=src_key_padding_mask)[0]

之后送到MultiheadAttention 的forward函数的attn_mask参数,而这里做的是一个self attention。

此时若是attn_mask出现一整行都是True的情况，通过如下源码中的实现mask的方法可以看到：

python 复制代码

    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_output_weights.masked_fill_(attn_mask, float('-inf'))
        else:
            attn_output_weights += attn_mask

把权重矩阵中需要mask的位置置为负无穷，然后再按行做softmax，问题就在这里，把一个元素全是是负无穷的tensor送给softmax，就会得到一个全是NaN的tensor。然后NaN和任何数运算都是NaN，NaN会传染，再经过一轮self attention，输出的tensor就全成NaN了。

解决方法：避免attention mask中存在全1/True的行