0
$\begingroup$

I am working on a seq-to-seq Image Captioning model, with Vision Transformer as the encoder and a LSTM based model as a decoder. The output from the encoder is given as the hidden state and cell state of the decoder LSTM. Each text is given SOS, EOS tags and padded with PAD tags. For predicting the LSTM is fed the previous words as input till EOS tag is emitted or till the maximum sentence length. Loss is calculated till the EOS tag only and PAD tags are not considered. How can I implement batches while training, as it is taking very long time for training with one example training at a time.

Seq-to-seq model :

class seq2seq(nn.Module):

  def __init__(self, encoder, decoder, max_length, device):
    super().__init__()
    self.encoder = encoder
    self.decoder = decoder
    self.max_length = max_length
    self.device = device
  
  def forward(self, x):

    encoder_hidden_state = self.encoder(x)

    hidden = encoder_hidden_state.reshape(1,1,-1)
    decoder_input = torch.tensor([[SOS_TOKEN]])
    decoder_hidden_state = (hidden.to(self.device), hidden.to(self.device))

    outputs = torch.zeros(self.max_length, x.shape[0], self.decoder.output_dim)


    for t in range(self.max_length):

      decoder_output, decoder_hidden_state = self.decoder(decoder_input.to(self.device), decoder_hidden_state)
      outputs[t] = decoder_output
      topv, topi = decoder_output.topk(1)
      decoder_input = topi.reshape(decoder_input.shape)
      if topi.item() == EOS_TOKEN:
        break
    
    return outputs
```
$\endgroup$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.