Splitting the Encoder and Decoder in an Autoencoder Architecture

I am working on replicating a dimensionality reduction model that, at its core, comprises a simple autoencoder. My naive implementation was as follows:

class encoder(nn.Module):
  def __init__(
          self,
          N,
          K
      ):
          super().__init__()
          self.U = nn.Linear(N,K)

  def forward(self,x):
    encoded_x = self.U(x)
    return encoded_x

class decoder(nn.Module):
  def __init__(
          self,
          N,
          K
      ):
          super().__init__()
          self.V = nn.Linear(K,N)

  def forward(self,x):
    decoded_x = self.V(x)
    return decoded_x

Creating a separate class for the encoder and decoder, and computing their inputs in the forward pass for the loss function:

encoded_X = encoder(X)
decoded_X = decoder(encoded_X)

loss = my_loss_func(X, encoded_X, decoded_X,...)

However, I have noticed that most autoencoders put the encoder and decoder together in a single class:

class SCA_autoencoder(nn.Module): 

    def __init__(
        self,
        N,
        K
    ):
        super().__init__()
        self.U = nn.Linear(N,K)
        self.V = nn.Linear(K,N)

    def forward(self,x):
        encoded_x = self.U(x)
        decoded_x = self.V(encoded_x)
        return encoded_x, decoded_x

When I test these two different implementations, the latter appears to converge while the former does not. Intuitively, this doesn’t make sense to me. Beyond readability, is there a difference in the way in which the gradients flow between these two different implementations? I can provide more code snippets or context if people want; the architecture/model I am trying to replicate is from this repo:

Albert

There shouldn’t be any principle difference. Having the encoder and decoder as own classes is very common as it yields cleaner coder, particularly once the encoder and decoder become more complex.

You could try to ensure that both version give you the same initial output for a batch. For this you create instances of both version and copy the weights and biases from one instance to the other; something like:

with torch.no_grad():
    layer2.weight.copy_(layer1.weight)
    layer2.bias.copy_(layer1.bias)

At least this gives you some sanity check by comparing both outputs and the loss. If those match, you can perform on backward pass to see if the gradients match.

Thank you for the response; I checked the outputs and they are identical.