How to Define Model in Native Pytorch Format and Import Into LightningModule Without Copy and Pasting?

[stanleyshly]

Let's say I have a native pytorch model like this

class NormalAutoEncoder(nn.Module)):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))

    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        embedding = self.encoder(x)
        return embedding

What is the best/most simple way to have the init function and forward function(basically, the whole network) into a pytorch lighting module without copy and pasting?

[rohitgr7]

just replace nn.Module -> nn.LightningModule.

[tchaton]

And implement the training_step to call forward directly. Don't forget to return the loss.

[cgerum]

Alternatively, you can just instantiate the Pytorch Module inside Lightning Module e.g.

class LightningBert(LightningModule):
     def __init__(self, **kwargs):
          self.autoencoder =  NormalAutoEncoder()

    def forward(self, x):
         return self.autoencoder(x)
         
    def train_step(self, ...):
          x = self.autoencoder(batch)
          loss = loss(x)
          return loss