How to convert CNN -> LSTM cascaded models to PyTorch Lightning

[amitpj]

Hi,

I am trying to covert PyTorch code for CNN -> LSTM cascaded models to PyTorch Lightning.
There are two nn.Module classes in the PyTorch code, one for CNN (encoder) and one for LSTM (decoder); the last hidden layer of CNN is used as input to LSTM. The architecture is described in this article: https://medium.com/@deepeshrishu09/automatic-image-captioning-with-pytorch-cf576c98d319

So, after conversion to PyTorch Lightning, there are two pl.LightningModule classes. I want to know how to populate the required methods in these two classes viz. forward(), training_step(), train_dataloader() and configure_optimizers().

Here are loss and optimizer definitions in PyTorch; optimizer uses parameters from both encoder and decoder models:
loss_criterion = nn.CrossEntropyLoss()
parameters = list(decoder_model.parameters()) + list(encoder_model.linear_layer.parameters()) + list(encoder_model.batch_norm.parameters())
optimizer = torch.optim.Adam(parameters, lr=0.001)

Here is the relevant PyTorch code from nested for training loop, encoder output is fed as input to the decoder and then the loss is calculated using the decoder output:
feats = encoder_model(imgs)
outputs = decoder_model(feats, caps, lens)
loss = loss_criterion(outputs, tgts)
decoder_model.zero_grad()
encoder_model.zero_grad()
loss.backward()
optimizer.step()

Regards,
Amit