Cut one part of an autoencoder (pytorch)

Question

I have a simple autoencoder architecture in PyTorch, which I train to do feature compression and reconstruction. My goal is to use the latent space of the autoencoder to reduce the initial dimensionality of my data and compress it in the test phase. To perform this, I would need to pass my test data only to my encoder, not the whole autoencoder. Would you have any idea of how to do this ? Something like model = Autoencoder.encoder() or else? My complete architecture is below:

class Autoencoder(nn.Module):
    def __init__(self, n_features):
        super(Autoencoder, self).__init__()
        
        self.n_sensors = n_sensors
        
        self.encoder = nn.Sequential(
            nn.Linear(self.n_features, 1),
            nn.ReLU(True))
        
        
        self.decoder = nn.Sequential(
            nn.Linear(1, self.n_features),
            nn.ReLU(True))

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

I think you can do (model = Autoencoder.encoder()). – Buoy Rina Sep 13 '21 at 09:36 — Buoy Rina, Sep 13 '21 at 09:36

score 0 · Answer 1 · answered Sep 13 '21 at 09:38

Base on your model definition, you can call forward on the encoder submodule directly:

class Autoencoder(nn.Module):
  def __init__(self, n_features):
      super(Autoencoder, self).__init__()
      
      self.n_features = n_features
      
      self.encoder = nn.Sequential(
          nn.Linear(self.n_features, 1),
          nn.ReLU(True))
      
      self.decoder = nn.Sequential(
          nn.Linear(1, self.n_features),
          nn.ReLU(True))

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

Keep in mind you first need to initialize the model:

>>> ae = Autoencoder(n_features=10)
>>> x = torch.empty(16, 10)
>>> ae.encoder(x).shape
(16, 1)

Thank you for your answer. In my case, I train the whole autoencoder (i.e. encoder + decoder) and I then want to use only the encoder in test mode, using the weights from my trained model. How can I load these weights to ```ae.encoder(x)``` ? — Martin, Sep 13 '21 at 09:53
If you have it already initialized, you can just call the `encoder` attribute, yes. — Ivan, Sep 13 '21 at 10:08

Payam Jome Yazdian · Answer 2 · 2023-03-30T18:34:48.747

First, regarding the model you provided, you can immediately call the model's encoder. However, sometimes it's not that easy and you might have complicated architecture including several layers such as linear layers, CCNs, RNNs, a Vector Quantizer, etc. In that case, I usually define two additional functions called forward_encoder and forward_decoder especially when I want to play with the latent space such as representation learning for clustering, visualization, and data compression.

class Autoencoder(nn.Module):
  def __init__(self, n_features):
      super(Autoencoder, self).__init__()
      
      self.n_features = n_features
      
      self.encoder = nn.Sequential(
          nn.Linear(self.n_features, 1),
          nn.ReLU(True))
      
      self.decoder = nn.Sequential(
          nn.Linear(1, self.n_features),
          nn.ReLU(True))

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

  def forward_encoder(self, x):
      latent = self.encoder(x)
      return latent

  def forward_decoder(self, latent):
      x = self.decoder(latent)
      return x

You could also modify the forward function and use a flag to whether return the latent variable or the reconstructed input.

  def forward(self, x, get_latent=Flase):
      latent  = self.encoder(x)
      x = self.decoder(latent)
      if get_latent:
          return latent:
      else:
      return x

Note that if you care about the latent space in the middle of the network rather than the output quality there are better options such as Denoising Autoencoder (DAE), Variational Autoencoder (VAE), Vector-Quantized Atuoencoder (VQ-VAE), etc. instead of Vanilla Autoencoder. Check this answer where I compared Vanilla Autoencoder and VAE. Note that there are other parameters that play role in the learned latent space such as parameters initialization, network architecture/complexity, activation functions, optimizer and so forth.

After completing the training, it is necessary to save the weights for future use (check this page for more details). Then, you need to create an instance of the model and load the saved weights into it for performing inference including the entire model or a specific component of it such as the encoder.

For example train the model:

# Create an instance of the model
model = Autoencoder(n_features=10)

# train the model
Train(model, epochs, dataloader, ...) 

# save the model parameters
torch.save(model.state_dict(), Save_Path)

At the inference time:

model = Autoencoder(n_features=10)
model.load_state_dict(torch.load(Save_Path))
#set the model to evaluation mode
model.eval()

sample_x = torch.zeros(16, 10)

out1 = model.encoder(sample_x)
out2 = model.forward_encoder(sample_x)
out3 = model(sample_x, get_latent=True)

where all out1, out2, and out3 would be the same.

Cut one part of an autoencoder (pytorch)

2 Answers2