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

user_14470826 0
Viewed 94

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

How can I have the __init__ and forward functions (basically, the whole network) into a pytorch lighting module without copying and pasting?

pytorch-lightning pytorch python
1 Answers

Simple. Make use of Python's inheritance mechanism.

If the following is a native PyTorch module

class NormalAutoEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = ...
        self.decoder = ...

    def forward(self, x):
        embedding = ...
        return embedding

Then have your new LightningAutoEncoder that also inherits from NormalAutoEncoder

class LightningAutoEncoder(LightningModule, NormalAutoEncoder):
    def __init__(self, ...):
        LightningModule.__init__(self) # only LightningModule's init
        NormalAutoEncoder.__init__(self, ...) # this basically executes __init__() of the NormalAutoEncoder

    def forward(self, x):
        # offloads its execution to NormalAutoEncoder's forward() method
        return NormalAutoEncoder.forward(self, x)

That's it. No copy-pasting.

user_51814640
Related