nn.TransformerEncoder - all nan values when src_key_padding_mask provided

Semir_Elezovikj: <p>I have a use case where I am dealing with sequences of variables lengths. I am looking to make use of <code>nn.TransformerEncoder</code>, but for some reason if a given sequence is of a length &lt; max_length of sequence, all values result in <code>nan</code> in the forward pass.</p>
<h2>
<a name="explanation-of-the-code-1" class="anchor" href="#explanation-of-the-code-1"></a><strong>Explanation of the code</strong>
</h2>
<p>I am including the fully reproducible code below.</p>
<p>In the example below, <code>batch_size</code> is 2, the maximum sequence length is 3, and I provide two sequences: one with 2 tokens and the other one with 3 tokens. Naturally, the sequence with 2 tokens needs to be padded in order to be fed to <code>nn.TransformerEncoder</code>.</p>
<p>In order to do this, I need to provide <code>src_key_padding_mask</code> of shape <code>(N, S)</code> where N is the <code>batch_size</code> and S is the <code>sequence_length</code>, in order to provide <strong>per-batch</strong> padding mask. For the sequence of length 2, a value of True needs to be provided for the position that is being padded - this means that for this particular sequence the padding values will be <code>[False, False,  True]</code></p>
<p>To my surprise, providing the <code>src_key_padding_mask</code> results with all <code>nan</code> values for the <strong>entire sequence</strong> where the sequence length &lt; max_sequence_length.</p>
<p><em>Note:</em> I have looked through all the other issues where other users were getting all <code>nan</code> values, but their scenarios were due to “user error”, where they provided sequences where all tokens are padded. This is not the case with my scenario.</p>
<pre><code class="lang-auto">from torch.nn.utils.rnn import pad_sequence
import torch
import torch.nn as nn
ego_encoder_layer = nn.TransformerEncoderLayer(d_model=2048, nhead=16, activation='gelu')
ego_transformer_encoder = nn.TransformerEncoder(ego_encoder_layer, num_layers=6)

sequences = []
batch_size = 2
src_padding_mask = torch.zeros((batch_size, 3)).type(torch.bool)

sequences.append(torch.rand(2, 2048))
sequences.append(torch.rand(3, 2048))
seq_lengths = [2, 3]

ego_seq2 = pad_sequence(sequences, padding_value=float('-inf'))
print(ego_seq2.shape)

ego_transformer_features = ego_transformer_encoder(ego_seq2, src_key_padding_mask=src_padding_mask)
print(ego_transformer_features.shape)




max_seq_length = 3
for ind, sl in enumerate(seq_lengths):
    for j in range(sl, max_seq_length):
        src_padding_mask[ind, j] = True # True means this is a padding token - skip

print('- seq_lengths, src_padding_mask -')
print(seq_lengths, src_padding_mask)

print('transformer features')
print(ego_transformer_features.shape)
ego_transformer_features[:, :, :5]
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">torch.Size([3, 2, 2048])
torch.Size([3, 2, 2048])
- seq_lengths, src_padding_mask -
[2, 3] tensor([[False, False,  True],
        [False, False, False]])
transformer features
torch.Size([3, 2, 2048])
tensor([[[    nan,     nan,     nan,     nan,     nan],
         [ 1.0203, -0.1215,  1.7757,  2.8037,  0.2745]],

        [[    nan,     nan,     nan,     nan,     nan],
         [ 1.4554,  0.1737,  1.9316,  2.8318,  1.0204]],

        [[    nan,     nan,     nan,     nan,     nan],
         [ 0.4414,  1.2467,  0.6188,  2.2753,  0.2954]]],
       grad_fn=&lt;SliceBackward&gt;)
</code></pre>

Semir_Elezovikj: <p><a class="mention" href="/u/alband">@albanD</a> Any thoughts on this?</p>

giladcohen: <p><a class="mention" href="/u/semir_elezovikj">@Semir_Elezovikj</a> I am experiencing the exact same NaN issue at the output of the transformerEncoder with padding. Did you find a solution for this issue?</p>

Mxttak: <p>For me,</p>
<pre><code class="lang-auto">padding_value=float('-inf')
</code></pre>
<p>was causing the same issue (actually I thought it were the <code>NaN</code>s in the tensors and I used <code>torch.nan_to_num(my_tensor, nan="-inf")</code>).</p>
<p>Replacing with</p>
<pre><code class="lang-auto">padding_value=0.
</code></pre>
<p>worked for me. Then, the mask could remain <code>bool</code></p>

makibrte: <p>what worked for me was setting the pad mask to float:</p>
<pre><code class="lang-auto">tgt_pad_mask = (tgt == 0).float()
</code></pre>