What is an idiomatic way to access the encoder layer in an autoencoder?

Lets say I have an autoencoder class

class Autoencoder(nn.Module):
  def __init__(self):
    self.encoder = nn.Sequential(...)
    self.decoder = nn.Sequential(...)

  def forward(self, x):
    x = self.encoder(x)
    x = self.decoder(x)
    return x

I have trained this model to my satisfaction. I now have some inputs I want to just encode. I have thought of two solutions:

  • adding a method that return self.encoder(x), or
  • separating this class into two classes

Are those approaches idiomatic?

If you have ae = Autoencoder() just using ae.encoder(x) is fine. Also enc = ae.encoder will give you the encoder module and it’ll work on it’s own, too (you can even save and load etc.).

Best regards

Thomas