Question about using another model in a customized loss function (grad None error))

BrokenArrow: <p>Hi, I would like to use model a to predict another model b(LeNet)'s parameters, so I need to implement a loss function myself</p>
<pre><code class="lang-auto">def custom_loss_function_new(output_batch, features, targets):
    batch_loss = 0.0
    features = features.to(DEVICE)
    targets = targets.to(DEVICE)

    for output in output_batch:
        output1 = output.clone()
        test_model = LeNet5(NUM_CLASSES, GRAYSCALE).to(DEVICE)
        optimizer1 = torch.optim.Adam(test_model.parameters(), lr=LEARNING_RATE1)  
        
        unflatten_parameters(output, test_model) # override LeNet's parameters
        test_model.train()

        logits, probas = test_model(features)
        cost1 = F.cross_entropy(logits, targets)
        cost1 = cost1.requires_grad_()
        print("COST1:", cost1)

        optimizer1.zero_grad()
        cost1.backward()

        for name, param in test_model.named_parameters():
                if param.requires_grad:
                    print(name, param.grad)

        optimizer1.step()
        loss = F.mse_loss(output1, flatten_parameters(test_model))
        # print("loss1:",loss)
        batch_loss += loss
    average_loss = batch_loss / len(output_batch)
    return average_loss
</code></pre>
<p>above is my loss function and it works well, how ever I want to define my own backforward function in the autograd, so I define a torch.autograd.Function class</p>
<pre><code class="lang-auto">class My_class_loss(torch.autograd.Function):
    @staticmethod
    def forward(ctx, output_batch, features, targets):
        average_loss = custom_loss_function_new(output_batch, features, targets)

        return average_loss
    
    @staticmethod
    def backward(ctx, grad_output):
        grads, = ctx.saved_tensors
        return grads, None, None
</code></pre>
<p>and use below code to calculate loss</p>
<pre><code class="lang-auto">        cost = custom_loss_function_new(generated_data, features, targets) #1
        cost = My_class_loss.apply(generated_data, features, targets) #2
        
</code></pre>
<p>when I use custom_loss_function_new, it just print the correct param.grad<br>
but when I use My_class_loss, I just got None for grad</p>
<pre><code class="lang-auto">COST1: tensor(3.3397, requires_grad=True)
features.0.weight None
features.0.bias None
features.3.weight None
features.3.bias None
classifier.0.weight None
classifier.0.bias None
classifier.2.weight None
classifier.2.bias None
classifier.4.weight None
classifier.4.bias None
</code></pre>
<p>I have no idea why I get different results in a same function</p>

KFrank: <p>Hi Broken!</p>
<aside class="quote no-group quote-modified" data-username="BrokenArrow" data-post="1" data-topic="197884" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/brokenarrow/48/68054_2.png" class="avatar"> BrokenArrow:</div>
<blockquote>
<pre><code class="lang-auto">class My_class_loss(torch.autograd.Function):
    @staticmethod
    def forward(ctx, output_batch, features, targets):
        average_loss = custom_loss_function_new(output_batch, features, targets)

        return average_loss
    
    @staticmethod
    def backward(ctx, grad_output):
        grads, = ctx.saved_tensors
        return grads, None, None
</code></pre>
<p>…<br>
when I use My_class_loss, I just got None for grad</p>
</blockquote>
</aside>
<p>Your <code>forward()</code> function doesn’t call <code>ctx.save_for_backward (...)</code>,<br>
so in your <code>backward()</code> function <code>ctx.saved_tensors</code> returns nothing<br>
(that is, an empty tuple).  Therefore your <code>grad</code>s come out as <code>None</code>.</p>
<p>Quoting from some of pytorch’s <a href="https://pytorch.org/docs/2.2/notes/extending.html" rel="noopener nofollow ugc">autograd documentation:</a></p>
<blockquote>
<p><strong>Step 2:</strong> It is your responsibility to use the functions in <code>ctx</code> properly in order to ensure that the new <a href="https://pytorch.org/docs/2.2/autograd.html#torch.autograd.Function" rel="noopener nofollow ugc"><code>Function</code></a> works properly with the autograd engine.</p>
</blockquote>
<ul>
<li><a href="https://pytorch.org/docs/2.2/generated/torch.autograd.function.FunctionCtx.save_for_backward.html#torch.autograd.function.FunctionCtx.save_for_backward" rel="noopener nofollow ugc"><code>save_for_backward()</code></a> must be used to save any tensors to be used in the backward pass.</li>
</ul>
<p>Best.</p>
<p>K. Frank</p>

BrokenArrow: <p>Thanks for your answer!<br>
I understand that I don’t have a complete implementation of the <code>backward</code> function in <code>My_class_loss</code>, which would cause me to not return the grad value correctly when I call this class to calculate the loss.<br>
But my problem is that I’m calling LeNet for inference in the <code>forward</code> function and my LeNet is not working correctly (LeNet’s grad is none), and I don’t think backward can affect LeNet.<br>
The strange thing for me is that if I call <code>custom_loss_function_new</code> directly, LeNet works fine, but if I call <code>custom_loss_function_new</code> in <code>My_class_loss</code>, it goes wrong, and I doubt if this has anything to do with <code>torch.autograd.Function</code> or <code>@staticmethod</code> implementation mechanism.</p>

BrokenArrow: <p>Hi KFrank,<br>
Here is the final version of <code>My_class_loss</code> for your reference.</p>
<pre><code class="lang-auto">
class My_class_loss(torch.autograd.Function):
    @staticmethod
    def forward(ctx, output_batch, features, targets):
        batch_loss = 0.0
        features = features.to(DEVICE)
        targets = targets.to(DEVICE)
        
        grads = torch.empty((len(output_batch), len(output_batch[0])), device=DEVICE)

        for i, output in enumerate(output_batch):
            output1 = output.clone()
            test_model = LeNet5(NUM_CLASSES, GRAYSCALE).to(DEVICE)
            optimizer1 = torch.optim.Adam(test_model.parameters(), lr=LEARNING_RATE)  
            test_model.train()
            unflatten_parameters(output1, test_model)

            logits, probas = test_model(features)
            cost1 = F.cross_entropy(logits, targets)
            cost1 = cost1.requires_grad_()

            optimizer1.zero_grad()
            cost1.backward()
            batch_loss += cost1

            for name, param in test_model.named_parameters():
                if param.requires_grad:
                    print(name, param.grad)

            grads[i] = torch.cat([p.grad.flatten() for p in test_model.parameters()])


        average_loss = batch_loss / len(output_batch)

        ctx.save_for_backward(grads)

        # print(output_batch)
        # print(features)
        # print(targets)

        # average_loss = custom_loss_function_new(output_batch, features, targets)

        return average_loss
    
    @staticmethod
    def backward(ctx, grad_output):
        grads, = ctx.saved_tensors
        return grads, None, None

</code></pre>
<p>so the problem here is that I got None for test_model(LeNet)'s grad<br>
the code in forward is just copied from <code>custom_loss_function_new</code><br>
and as I said, if I call <code>custom_loss_function_new</code> directly, everything is fine, i will get the correct grad</p>
<p>and below is my training code just in case you need to take a look</p>
<pre><code class="lang-auto">for epoch in range(NUM_EPOCHS):
    
    model.train()
    for batch_idx, (features, targets) in enumerate(train_loader):
        features, targets = features.to(DEVICE), targets.to(DEVICE)
        
        z = torch.rand((BATCH_SIZE, NUM_INPUT)).to(DEVICE)
        
        ### FORWARD AND BACK PROP
        generated_data = model(z)
        
        cost = My_class_loss.apply(generated_data, features, targets)
        # cost = custom_loss_function_new(generated_data, features, targets)

        optimizer.zero_grad()
        cost.backward()

        optimizer.step()


        ### LOGGING
        # if not batch_idx % 50:
        #     print('Epoch: %03d/%03d | Batch %04d/%04d | Cost: %.4f' 
        #           % (epoch+1, NUM_EPOCHS, batch_idx, len(train_loader), cost))
        print('Epoch: %03d/%03d | Batch %04d/%04d | Cost: %.4f' 
                  % (epoch+1, NUM_EPOCHS, batch_idx, len(train_loader), cost))
        


        model.eval()
        with torch.set_grad_enabled(False): # save memory during inference
            z = torch.rand(1, NUM_INPUT).to(DEVICE)
            samples = model(z)

            test_model = LeNet5(NUM_CLASSES, GRAYSCALE).to(DEVICE)

            samples = samples.squeeze()
            unflatten_parameters(samples, test_model)

            print("accuracy:", compute_accuracy(test_model, test_loader).item())
            
        print('Time elapsed: %.2f min' % ((time.time() - start_time)/60))

</code></pre>

KFrank: <p>Hi Broken!</p>
<aside class="quote no-group" data-username="BrokenArrow" data-post="3" data-topic="197884" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/brokenarrow/48/68054_2.png" class="avatar"> BrokenArrow:</div>
<blockquote>
<p>But my problem is that I’m calling LeNet for inference in the <code>forward</code> function and my LeNet is not working correctly (LeNet’s grad is none),</p>
</blockquote>
</aside>
<p>My apologies – I misunderstood your question.</p>
<p>What’s going on is that autograd doesn’t track computations inside of an<br>
<code>autograd.Function</code>.  It’s as if the computations are performed inside of<br>
a <code>with torch.no_grad():</code> block.</p>
<p>Here is a simple script that illustrates this:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

x = torch.tensor ([1.0, 2.0], requires_grad = True)

def someFunction (x):
    print ('torch.is_grad_enabled() =', torch.is_grad_enabled())
    y = 2.0 * x
    print ('y =', y, '(inside someFunction())')   # tracked by autograd
    return  y

u = someFunction (x)
print ('u =', u)                                  # tracked by autograd

class SomeCustomFunction (torch.autograd.Function):
    @staticmethod
    def forward (ctx, x):
        z = someFunction (x)                      # no longer tracked by autograd
        return  z
    
    @staticmethod
    def backward(ctx, grad_output):
        print ('ctx.saved_tensors = ...')
        return  None

custom_fn = SomeCustomFunction()
u = custom_fn.apply (x)
print ('u =', u)                                  # use of SomeCustomFunction is tracked by autograd, but not its internals
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.2.1
torch.is_grad_enabled() = True
y = tensor([2., 4.], grad_fn=&lt;MulBackward0&gt;) (inside someFunction())
u = tensor([2., 4.], grad_fn=&lt;MulBackward0&gt;)
torch.is_grad_enabled() = False
y = tensor([2., 4.]) (inside someFunction())
u = tensor([2., 4.], grad_fn=&lt;SomeCustomFunctionBackward&gt;)
</code></pre>
<p>I don’t know if there is a way to reliably reenable autograd tracking inside of a<br>
custom <code>Function</code>.  Is there some way you can recast your program logic so<br>
that it doesn’t use a custom <code>Function</code>?</p>
<p>Also, your <code>My_class_loss</code> takes <code>output_batch</code>, <code>features</code>. and <code>targets</code><br>
as arguments, but its <code>backward()</code> appears to be returning the gradients<br>
with respect to the <code>Parameter</code>s of a temporarily-instantiated <code>LeNet5</code>.  I see<br>
no reason that these gradients will match up with the input arguments, so I<br>
think a backward pass through <code>My_class_loss</code> will fail.</p>
<p>Last, depending on what you actually want to compute the gradients<br>
with respect to, you might be wanting to backpropagate through the<br>
<code>optimizer1.step()</code> call (that seems to be missing from your “final<br>
version”).</p>
<p>If so, some of the discussion in the following thread might be relevant:</p>
<aside class="quote quote-modified" data-post="1" data-topic="196075">
  <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/9de0a6/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/backpropagate-model-based-on-performance-of-another-model/196075">Backpropagate model based on performance of another model</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Hi, 
I currently have a problem where I have two models, model_1 and model_2, and I would like to optimise one based on the performance of the other. 

Both models are of different types.
I would like to optimise model_2 based on the performance of model_1. The loss function for model_1 is well defined (CrossEntropyLoss), but for model_2 it is not.
model_2 produces an output, which implicitly effects the performance of model_1. The output of model_2 is not directly given to model_1. Overall, it’…
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>

BrokenArrow: <p>Hi KFrank<br>
Thanks a lot for your help, I successfully solve the problem: just add <code>with torch.enable_grad():</code> in the <code>forward</code> function. And that’s all.<br>
<img src="https://discuss.pytorch.org/images/emoji/apple/partying_face.png?v=12" title=":partying_face:" class="emoji only-emoji" alt=":partying_face:" loading="lazy" width="20" height="20"></p>