Using decoder as part of loss function

mtnpickles: <p>Hi,</p>
<p>I have trained an autoencoder, e.g.:</p>
<pre><code class="lang-auto">class autoencoder(nn.Module):
    def __init__(self):
        super().__init__()

        self.encoder = nn.Sequential( )
        self.decoder= nn.Sequential( )

    def forward(self,x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x
</code></pre>
<p>and let’s say</p>
<pre><code class="lang-auto">AE_model = autoencoder() 
</code></pre>
<p>While I’m training another new model e.g.:</p>
<pre><code class="lang-auto">new_model = NewModel()
</code></pre>
<p>I tried to use the decoder part of the autoencoder as part of my loss function. First, I freeze the parameters:</p>
<pre><code class="lang-auto">for param in AE_model.parameters():
   param.requires_grad=False
</code></pre>
<p>Then, during training:</p>
<pre><code class="lang-auto">mseloss = nn.MSELoss()
optimizer = torch.optim.Adam(new_model.parameters())
</code></pre>
<pre><code class="lang-auto">prediction = new_model(batch_x)
decoded_prediction = AE_model.decoder(prediction)
Loss = mseloss(decoded_prediction, batch_y)
Loss.backward()
optimizer.step()
</code></pre>
<p>Problem: the loss doesn’t seem to decrease and is always stuck at one value… Can someone please help?</p>
<p>Thanks.</p>

InnovArul: <aside class="quote no-group" data-username="mtnpickles" data-post="1" data-topic="150553">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/958977/48.png" class="avatar"> mtnpickles:</div>
<blockquote>
<p><code>decoded_prediction = AE_model.decoder(prediction)</code></p>
</blockquote>
</aside>
<p>I am not sure of the exact reason.<br>
Are you setting <code>AE_model</code> to eval mode?</p>

mtnpickles: <p>I did… Also, the <code>new_model</code> is training fine when I didn’t use the <code>AE_model.decoder</code>, i.e., when I just predict the encoded outputs directly.</p>

InnovArul: <aside class="quote no-group" data-username="mtnpickles" data-post="3" data-topic="150553">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/958977/48.png" class="avatar"> mtnpickles:</div>
<blockquote>
<p>when I just predict the encoded outputs directly.</p>
</blockquote>
</aside>
<p>Can you elaborate a bit more on this and the motivation for doing this?<br>
What are the encoded outputs and how do you predict them directly?</p>

ksmdanl: <p>For debugging purpose, I would to try to remove the encoder part in <code>autoencoder()</code> and attach <code>NewModel()</code> to <code>autoencoder()</code> as the new encoder, train only the new attached encoder part while the trained decoder stays freezed, as you did here, and see what changes.</p>

mtnpickles: <p>So I trained an autoencoder (<code>AE_model</code>) on a dataset of Domain B. I want to train another model (<code>new_model</code>) that takes in dataset of Domain A, and output predictions in Domain B.</p>
<p>First, I encode the dataset of Domain B using (<code>AE_model.encoder</code>).</p>
<p>Then, my idea was to have <code>new_model</code> to predict encoded dataset of Domain B, then have <code>AE_model.decoder</code> to decode back to the original state.</p>
<p>When I tried using a Loss function of:</p>
<pre><code class="lang-auto">encoded_prediction = new_model(X_valid)
Loss(encoded_prediction, encoded_GT)
</code></pre>
<p>it works fine. The training loss decreases nicely and so does my validation loss.</p>
<p>Then, when I tried incorporating the decoder as part of my Loss function:</p>
<pre><code class="lang-auto">encoded_prediction = new_model(X_valid)
decoded_prediction = AE_model.decoder(encoded_prediction )
Loss(decoded_prediction , decoded_GT)
</code></pre>
<p>the training and validation loss do not decrease at all.</p>

mtnpickles: <p>I did exactly just this:</p>
<pre><code class="lang-auto">class combined(nn.Module): 
    def __init__(self):
        super().__init__()
        self.AE_model = autoencoder()
        self.AE_model.load_state_dict(torch.load('some_weights.pth'))
        
        #####Freezing all weights in AE
        for param in self.AE_model.parameters():
            param.requires_grad = False

        #### New Model
        self.block = nn.Sequential ( )
       
    def forward(self, x):
        x = self.block(x)
        x = self.AE_model.decoder(x)
        return x
</code></pre>
<p>The training and validation loss still does not change a bit. <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"> . This is so weird, as when I do the Loss directly at the encoded prediction, it works fine. Just by adding the decoder made it fail. <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"></p>

mtnpickles: <p>Update: I may have found the problem and a solution. I think it turns out that my <code>encoded_prediction</code> has positive values while my <code>new_model</code> takes in positive and negative values. So I guess the problem is very  nonlinear and the NN wasn’t able to initialize the weights properly in the <code>encoded_prediction</code>, thus the <code>autoencoder.decoder</code> is having a hard time finding a gradient of steepest descent. I fixed it by changing my ReLU to Tanh activations.</p>