Setting custom gradient for backward

EvilSKy: <p>Hi,</p>
<p>Normally one would use a forward pass to calculate a loss and then perform a backward pass and the gradient is automatically calculated.</p>
<p>My situation is that I don’t have the loss but I have the gradient calculated. My question is how do I set a custom gradient to a network (fully connected) and run the backward optimization using the custom gradient? Thanks</p>

ptrblck: <p>You could assign the precomputed gradients to the <code>.grad</code> attribute of all parameters and call <code>optimizer.step()</code> afterwards.  Here is a small example:</p>
<pre><code class="lang-python">lin = nn.Linear(1, 1, bias=False)
optimizer = torch.optim.SGD(lin.parameters(), lr=1e-3)

print(lin.weight)
lin.weight.grad = torch.ones_like(lin.weight) * 10
print(lin.weight.grad)
optimizer.step()
print(lin.weight)
</code></pre>

EvilSKy: <p>Thank you! What if I have a network? Like a two layers network. Can I assign grad to the last layer and then apply backward() to perform the back propogation for the first layer?</p>

ptrblck: <p>Yes, you could manipulate the gradients manually of any layer using the previous code snippet or you could also use hooks (via <code>register_hook</code>), which might be cleaner especially if you are using multiple layers.</p>

EvilSKy: <p>Thank you for your answer.<br>
I realized that I have a slightly different problem here.</p>
<p>Rather than changing the weight of the layer, what I would like to do is to set the gradient of the loss w.r.t the output of my network. So dLoss/dOutput, which is not part of the network layer. It is more like setting a grad for a tensor variable? To be more specific, my pseudo code would be:</p>
<pre><code class="lang-auto">output = myNetwork(input)  

my_grad = getMyGradient()  #my custom dLoss/dOutput

output.grad = my_grad
output.backward() #usually backward starts from loss, but I would like to start from the output.
optimizer.step()

</code></pre>
<p>What would be the correct implementation for this kind of task?<br>
Thank you in advance.</p>

ptrblck: <p>You can pass the gradient directly to the <code>backward</code> operation as an argument: <code>output.backward(my_grad)</code>.</p>

EvilSKy: <p>It is working! Thank you!</p>

medvisioner: <p>Hi, ptrblck</p>
<p>Could you give an example of using hooks to manipulate the gradient of each layer (say a three-layer fully connected network)? I am new to PyTorch and never use hooks before.</p>

ptrblck: <p>Sure, here is a small example:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 2)
        self.fc2 = nn.Linear(2, 2)
        self.fc3 = nn.Linear(2, 1)
        
    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = F.relu(x)
        x = self.fc3(x)
        return x
    
model = MyModel()
x = torch.randn(1, 1)

model.fc1.weight.register_hook(lambda grad: grad+1000.)
model.fc2.bias.register_hook(lambda grad: grad-1000.)
model.fc3.weight.register_hook(lambda grad: grad*1000.)

out = model(x)
out.mean().backward()
print([(name, param.grad) for name, param in model.named_parameters()])
# [('fc1.weight', tensor([[ 999.7733],
#         [1000.0000]])), ('fc1.bias', tensor([-0.2406,  0.0000])), ('fc2.weight', tensor([[0.0000, 0.0000],
#         [0.3678, 0.0000]])), ('fc2.bias', tensor([-1000.0000,  -999.4205])), ('fc3.weight', tensor([[ 0.0000, 67.0654]])), ('fc3.bias', tensor([1.]))]
</code></pre>

wcy8902: <p>Hi, ptrblck.<br>
I want to train a model, its forward path includes three part A, B and C,sequentially.<br>
The main problem I encountered is that B is not differentiable,<br>
so I replace it with a simulation model B_prime in backward part for gradient computation.<br>
(But still use original B in forward path)<br>
Any suggestion regarding how to build a custom backpropagation structure?<br>
Thanks!</p>

ptrblck: <p>If you are using non-differentiable operations and want to implement a custom <code>backward</code> you could defined a custom <code>autograd.Function</code> as described <a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html">here</a>.</p>

wcy8902: <p>Hi, thanks for your reply.<br>
I’m using autograd.Function with nn.module now, but I get this <a href="https://discuss.pytorch.org/t/element-0-of-tensors-does-not-require-grad-and-does-not-have-a-grad-fn-how-can-i-fix-this/140756/2">error</a> when calling loss.backward() for mySys.<br>
I’ve seen your reply below and found my custom backward function’s output sim_out is detached from the computation graph(.grad_fn=None).<br>
Here is a snippet of my code:</p>
<pre><code class="lang-auto">class B_nd(torch.autograd.Function):
    def forward(ctx, input):
        ctx.save_for_backward(input)
        # some non-differentiable functions
        return out

    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        input.requires_grad_(True)
        sim_model = UNet()
        sim_model.load_state_dict(torch.load('/Desktop/CodeFolder/2023/pth/800.pth', map_location='cuda:0'))
        sim_model.to(torch.device("cuda"))
        sim_model.zero_grad()
        sim_out = sim_model(input)
        print(sim_out.grad_fn)
        sim_out.backward(gradient=grad_output)
        gradient = input.grad
        return gradient

class mySys(nn.Module):
    def __init__(self):
        super(mySys, self).__init__()
        self.A = nn.Conv2d(3, 3, kernel_size=1)
        self.C = nn.Conv2d(3, 3, kernel_size=1)
    
    def forward(self, x):
        a_out = self.A(x)
        b_out = B_nd.apply(a_out)
        c_out = self.C(b_out)
        return c_out
</code></pre>
<p>Any method or suggestion to attach the UNet in the backward part of B to the computation graph?</p>

ptrblck: <p>I don’t fully understand the explanation as nothing is detached in your code.<br>
You might just need to enable gradient calculation in your <code>backward</code> via <code>torch.set_grad_enabled(True)</code>.</p>