Train a Neural Network and using it as a loss function

a-parida12: <p>I am training a model and also want to use it as a loss function too. See schematic below:</p>
<pre><code class="lang-auto">pred = self.forward(im1, im2)

# calculate loss
with torch.no_grad():
     pred_new = self.forward(im3, pred)
loss = F.mse_loss(pred, pred_new)
</code></pre>
<p>When I do a <code>loss.backward()</code>, I get an RunTime Error as the <code>grad_fn</code> is missing.</p>
<p>What would be the ideal way to calculate the loss in this scenario?</p>

ptrblck: <p>Your use case should generally work as the target usually doesn’t require a gradient as seen here:</p>
<pre><code class="lang-python">model = nn.Linear(10, 2)

x = torch.randn(1, 10)
pred = model(x)

with torch.no_grad():
    pred_new = model(torch.randn(1, 10))
    
loss = F.mse_loss(pred, pred_new)
loss.backward()
print(model.weight.grad)
</code></pre>
<p>With that being said, I don’t know how your model is defined etc. and which part of the code is raising the error.</p>

a-parida12: <p>The difference between the example you show is that, in my second call to the <code>model</code> the prediction (model output) from the first call is used as an input without using <code>detach</code> .</p>

ptrblck: <p>Could you post a minimal, executable code snippet which would reproduce the issue, please?</p>

a-parida12: <p>This is exactly what I am doing and trying to achieve. You should be able to reproduce the <code>RuntimeError</code>.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.inp1 = nn.Linear(10, 2)
        self.inp2 = nn.Linear(10,2)
        self.out = nn.Linear(4, 10)

    def forward(self, im1, im2):
        im1 = self.inp1(im1)
        im2 = self.inp2(im2)
        return self.out(torch.concat([im1, im2], axis=1))

model = Model()
inp1 = torch.randn(1, 10)
inp2 = torch.randn(1, 10)
gt = torch.randn(1, 10)

pred = model(inp1, inp2)

with torch.no_grad():
    pred_new = model(pred, inp1)
    
loss = F.mse_loss(gt, pred_new)
loss.backward()
print(model.weight.grad)
</code></pre>
<p>The error is</p>
<blockquote>
<p>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</p>
</blockquote>

ptrblck: <p>Thanks for the code! The error is expected as none of the tensors is attached to a computation graph:</p>
<pre><code class="lang-python">loss = F.mse_loss(gt, pred_new)
</code></pre>
<p><code>gt</code> is the target tensor created directly via <code>gt = torch.randn(1, 10)</code> while <code>pred_new</code> is created in the <code>no_grad()</code> context so also not attached to a computation graph.<br>
Could you explain which gradients the <code>backward()</code> call should calculate?</p>

a-parida12: <p>I have <code>grad_fn</code> for the operations from <code>inp1</code> and <code>inp2</code> to the output <code>pred</code>. I want the gradients to computed for that part of the code but with respect to the loss calculated with the mse_loss.</p>
<p>so the parameters of the <code>model</code> can be updated with the step function of the optimizer.</p>

a-parida12: <p>I feel one way I could achieve this is  by writing the code like this. Could you verify the effect of the deepcopy from copy? What happens when the model is loaded on the GPU?</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import SGD
import copy

class Model(nn.Module):
    def __init__(self, grayscale=False):
        super().__init__()

        self.inp1 = nn.Linear(10, 2)
        self.inp2 = nn.Linear(10,2)
        self.out = nn.Linear(4, 10)

    def forward(self, im1, im2):
        im1 = self.inp1(im1)
        im2 = self.inp2(im2)
        return self.out(torch.concat([im1, im2], axis=1))

model = Model()
optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)
loss_model = copy.deepcopy(model)
inp1 = torch.randn(1, 10)
inp2 = torch.randn(1, 10)
gt = torch.randn(1, 10)

pred = model(inp1, inp2)
loss_model = copy.deepcopy(model)
pred_new = loss_model(pred, inp1)
loss = F.mse_loss(inp1, pred_new)
loss.backward()
optimizer.step()
</code></pre>
<p>Also the memory footprint of the training is super high due to the copy of the model with all its parameters. Is their an elegant method to prevent the copy of the model to achieve the same results?</p>

a-parida12: <p>I think in terms of saving on some memory at forward after copying I set the <code>requires_grad= False</code>.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import SGD
import copy

class Model(nn.Module):
    def __init__(self, grayscale=False):
        super().__init__()

        self.inp1 = nn.Linear(10, 2)
        self.inp2 = nn.Linear(10,2)
        self.out = nn.Linear(4, 10)

    def forward(self, im1, im2):
        im1 = self.inp1(im1)
        im2 = self.inp2(im2)
        return self.out(torch.concat([im1, im2], axis=1))

model = Model()
optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)

# copy the og model remove gradient calculation
loss_model = copy.deepcopy(model)
for p in loss_model.parameters():
        p.requires_grad = False

inp1 = torch.randn(1, 10)
inp2 = torch.randn(1, 10)
gt = torch.randn(1, 10)

pred = model(inp1, inp2)
pred_new = loss_model(pred, inp1)
loss = F.mse_loss(inp1, pred_new)
loss.backward()
optimizer.step()
</code></pre>