How to use nn.TransformerDecoder() at inference time

agnes: <p>Hello. I am using <code>nn.TransformerDecoder()</code> module to train a language model. During training time, the model is using target <code>tgt</code> and <code>tgt_mask</code>, so at each step the decoder is using the last true labels.<br>
However, for text generation (at inference time), the model shouldn’t be using the true labels, but the ones he predicted in the last steps. Can we do that with <code>nn.TransformerDecoder()</code> ? Or should I reimplement the module to add the new predicted labels at each step ?</p>

tomb88: <p>Were you able to get an answer on this? Is there an example?</p>

zhangguanheng66: <p>An example with <code>nn.Transformer</code> module for word language model <a href="https://github.com/pytorch/examples/tree/master/word_language_model" rel="nofollow noopener">here</a>.<br>
Or a tutorial <a href="https://pytorch.org/tutorials/beginner/transformer_tutorial.html" rel="nofollow noopener">here</a></p>

tomb88: <p>The example you posted is for the TransformerEncoder. I’m looking for an explicit usage example of the TransformerDecoder.<br>
Can you provide one?</p>

zhangguanheng66: <p>So for TransformDecoder, what’s your target and memory sequences?</p>

agnes: <p>Yes. You need to initialized the target with a tensor of SOS (start of sentence). At each step, you append this tensor with the pred, as explained in the paper <a href="https://arxiv.org/pdf/1706.03762.pdf" rel="noopener nofollow ugc">Attention Is All You Need</a>:</p>
<blockquote>
<p>At each step the model is auto-regressive[10], consuming the previously generated symbols as additional input when generating the next</p>
</blockquote>
<pre><code class="lang-auto">import torch
import torch.nn as nn

device = 'cpu'
d_model = 768
bs = 4
sos_idx = 0
vocab_size = 10000
input_len = 10
output_len = 12

# Define the model
encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=4).to(device)

encoder = nn.TransformerEncoder(encoder_layer=encoder_layer,
                                num_layers=6).to(device)

decoder_layer = nn.TransformerDecoderLayer(d_model=d_model, nhead=4).to(device)

decoder = nn.TransformerDecoder(decoder_layer=decoder_layer,
                                num_layers=6).to(device)

decoder_emb = nn.Embedding(vocab_size, d_model)

predictor = nn.Linear(d_model, vocab_size)

# for a single batch x
x = torch.randn(bs, input_len, d_model).to(device)
encoder_output = encoder(x)  # (bs, input_len, d_model)

# initialized the input of the decoder with sos_idx (start of sentence token idx)
output = torch.ones(bs, output_len).long().to(device)*sos_idx
for t in range(1, output_len):
    tgt_emb = decoder_emb(output[:, :t]).transpose(0, 1)
    tgt_mask = torch.nn.Transformer().generate_square_subsequent_mask(
        t).to(device).transpose(0, 1)
    decoder_output = decoder(tgt=tgt_emb,
                             memory=encoder_output,
                             tgt_mask=tgt_mask)

    pred_proba_t = predictor(decoder_output)[-1, :, :]
    output_t = pred_proba_t.data.topk(1)[1].squeeze()
    output[:, t] = output_t

#output (bs, output_len)

</code></pre>
<p>My final model looks like this:</p>
<pre><code class="lang-auto">class Transformer(nn.Module):
    def __init__(self):
        pass
    
    def forward(self, x, y=None, mode='train'):
        encoder_output = self.encoder(x)

        if mode == 'train':
            tgt_emb = self.decoder_emb(y).transpose(0, 1)
            tgt_mask = torch.nn.Transformer().generate_square_subsequent_mask(
                tgt_emb.size(0)).to(self.device).transpose(0,1)
            decoder_output = self.decoder(tgt=tgt_emb,
                                          tgt_mask=tgt_mask,
                                          memory=encoder_output)
            return self.predictor(decoder_output).permute(1, 2, 0)
        elif mode == 'generate':
            # the solution I gave before
            return 
        

</code></pre>

xdwang0726: <p>Do you need <code>tgt_mask</code> in the inference? I am wondering is the for loop in inference processing do the same job as subsquent mask? Also, I am wondering have you done a experiment with <code>nn.TransformerDecoder</code>? I have some trouble using it, as when inference, it always return the same token. I really appreciate if you can help! Thank you!</p>

Jingles: <p>I am learning seq2seq by working on a machine translation problem, from DE to EN.</p>
<p>Train and validation worked well. I got acceptable losses, and printing out the source, target and predicted sentence is acceptable too.</p>
<p>But for test (generate), I appended previously generated symbols as additional inputs for generating the next symbol. But my results output the same symbol again and again.</p>
<p>For example, the output I got is:</p>
<pre><code class="lang-auto">&lt;sos&gt; themselves themselves themselves themselves themselves themselves themselves themselves themselves 
</code></pre>
<p>Any idea what might have gone wrong?</p>
<p>Example output after every generated symbol:</p>
<pre><code class="lang-auto">tensor([2])
tensor([  2, 700])
tensor([  2, 700, 700])
tensor([  2, 700, 700, 700])
tensor([  2, 700, 700, 700, 700])
tensor([  2, 700, 700, 700, 700, 700])
tensor([  2, 700, 700, 700, 700, 700, 700])
tensor([  2, 700, 700, 700, 700, 700, 700, 700])
tensor([  2, 700, 700, 700, 700, 700, 700, 700, 700])
tensor([  2, 700, 700, 700, 700, 700, 700, 700, 700, 700])
tensor([  2, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700])
tensor([  2, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700])
tensor([  2, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700])
tensor([  2, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700, 700])
</code></pre>

mathematicsofpaul: <p>I suspect you wouldn’t need a <code>tgt_mask</code> during inference since there would be nothing to mask. Only a shifted <code>SOS</code> token.</p>

mathematicsofpaul: <p>Can post this code somewhere where I could see it?</p>

mathematicsofpaul: <p>Did you happen to solve this issue? I am getting the same repeated values as well… Strange since I made sure that I accounted for the required shift in the output for the decoder…</p>

asdf11: <p>Hi Im facing the same problem. How did you ensure the shift? Im not sure how to do it: Im using padded sequences as target input: e.g. <code>target_tensor = [1,2,3,4,5,99,0,0,0,0,0,0]</code> (<code>1</code> is sos-token, <code>99</code> is eos-token, cutting the last element doesnt make sense here, since it s a padding token.) I tried to switch the token with<code> target_tensor[target_tensor == 99] = 0</code> whcih results in an grading error.</p>
<p>Also maybe <a class="mention" href="/u/agnes">@Agnes</a> could you provide more parts of the code Im not sure where to initialize the code? In <strong>init</strong>? or outside the model?</p>
<p>Looking at this <a href="https://colab.research.google.com/drive/1g4ZFCGegOmD-xXL-Ggu7K5LVoJeXYJ75#scrollTo=-DtquWxxKF9I" rel="nofollow noopener">example</a>  none of the things <a class="mention" href="/u/agnes">@agnes</a> wrote seems necessary? Is there any explanation for that?</p>
<p>Adding this <a href="https://github.com/pytorch/tutorials/issues/719#issuecomment-612043938" rel="nofollow noopener">discussion</a> which descibes a similar apporach to <a class="mention" href="/u/agnes">@agnes</a></p>

mathematicsofpaul: <p><a class="mention" href="/u/asdf11">@asdf11</a> Can you check if your training outputs give values that do not repeat? I found that that was an issue I was facing. This issue of repeated training outputs is still an issue for me and I suspect that ‘tgt_mask’ doesnt really work <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"></p>

raullalves: <p>Does anyone have a code snippet about using the TransformerDecoder properly?</p>
<p>Regards</p>

david.waterworth: <p><a class="mention" href="/u/mathematicsofpaul">@mathematicsofpaul</a> - one thing I just noticed after banging my head against this for the last day - I removed the PositionalEncoding layer (from <a href="https://pytorch.org/tutorials/beginner/transformer_tutorial.html" rel="nofollow noopener">https://pytorch.org/tutorials/beginner/transformer_tutorial.html</a>) and that “fixed” the problem. Before I did this it simply repeated the same token (both training and validation).</p>
<p>Not sure if there’s a bug in it, I suspect that it’s necesary for a real transformer (I’m just trying to get it to work with a couple of training cases and using the same cases for validation) so I guess I’ll have to look closer at what it’s doing (and what it’s supposed to do)</p>

mathematicsofpaul: <p><a class="mention" href="/u/david.waterworth">@david.waterworth</a> thanks for posting! Woah, I believe the Positional Encoding module is a must since it actually gives order to the words you input into the model. If you included a sentence without Positional Encoding, that basically means you’re giving the model a bag of words instead of a sequence of words. Regardless, it might be that there is a bug in the Positional Encoding - maybe. Could you please post your code in a google colab codebook or simply post your model here?</p>

david.waterworth: <p>Yeah <a class="mention" href="/u/mathematicsofpaul">@mathematicsofpaul</a> , I doubt my model will work for a real dataset without it. This is how I implemented my forward function. I’m going to spend tomorrow looking more closely at the PositionalEncoding</p>
<pre><code>def forward(self, src, tgt):

    # create padding masks
    src_key_padding_mask = (src == self.src_vocab.stoi['&lt;pad&gt;'])  # N S
    tgt_key_padding_mask = (tgt == self.tgt_vocab.stoi['&lt;pad&gt;'])  # N T
    memory_key_padding_mask = src_key_padding_mask.clone()        # N S

    # prevent attention of future output tokens
    tgt_mask = self.transformer.generate_square_subsequent_mask(tgt.shape[-1])  # T T

    src = rearrange(src, 'n s -&gt; s n')
    tgt = rearrange(tgt, 'n t -&gt; t n')

    src = self.embed_src(src) * math.sqrt(self.d_model) # S, N, E
    tgt = self.embed_tgt(tgt) * math.sqrt(self.d_model) # T, N, E

    output = self.transformer(
        src, tgt, 
        tgt_mask=tgt_mask, 
        src_key_padding_mask=src_key_padding_mask,
        tgt_key_padding_mask=tgt_key_padding_mask, 

    output = rearrange(output, 't n e -&gt; n t e')
    return self.fc(output)
</code></pre>
<p>and my train and validation are the same</p>
<pre><code>        # reshape inputs
        src = rearrange(batch.object_name, 's n -&gt; n s')
        tgt = rearrange(batch.tags, 't n -&gt; n t')   

        # Create tgt_inp and tgt_out (which is tgt_inp but shifted by 1)
        tgt_inp, tgt_out = tgt[:,:-1], tgt[:,1:]

        optim.zero_grad()
        out = model(src, tgt_inp)
        #predicted = F.log_softmax(out, dim=-1).argmax(dim=-1)
        loss = criterion(rearrange(out, 'b t v -&gt; (b t) v'), rearrange(tgt_out, 'b o -&gt; (b o)'))

        # Backpropagate and update optim
        loss.backward()
        optim.step()</code></pre>

mathematicsofpaul: <p><a class="mention" href="/u/david.waterworth">@david.waterworth</a> What did you set for your <code>self.transformer</code>? I assume <code>nn.Transformer()</code>. Great let me know about the Positional Encoding! I too have been banging my head about this.</p>
<p>Also, check out this very informative <a href="https://github.com/pytorch/tutorials/issues/719" rel="nofollow noopener">thread</a>.</p>

david.waterworth: <p>Yes it’s nn.Transformer()</p>

david.waterworth: <p><a class="mention" href="/u/mathematicsofpaul">@mathematicsofpaul</a>  I don’t think the problem was with the position encoding. I think it was the initialisation. I’d copied the code below and applied it to the full model - it appears to mess something up  - when I removed it and added the position encoding back training converges to a much lower loss and the decoder stops repeating the same token over and over.</p>
<pre><code class="lang-auto">for p in model.parameters():
    if p.dim() &gt; 1:
        torch.nn.init.xavier_normal_(p)
</code></pre>
<p>Perhaps the code overwrites the position encoding as it’s stored in a buffer - although the docstring for ‘register_buffer’ does say it’s for model state that aren’t parameters? The other option is the embedding weights are initialised to zero - perhaps overwriting them with non-zeros messes with that layer?</p>