Backward() in custom layer is not called

hayden: <p>What I am going to do is modifying weight in Conv2d after loss.backward() and before optimizer.step(). One solution is to modify weights in corresponding layers after loss.backward(). I just want to make a custom layer to do this keeping train() function clean. Here is the code snippet:</p>
<pre><code class="lang-auto">class CustomConv(nn.Conv2d):
    def __init__(self, *kargs, **kwargs):
        super(CustomConv, self).__init__(*kargs, **kwargs)

    def forward(self, input):
        input = WeightModifier(self.weight)(input)
        out = nn.functional.conv2d(
            input,
            self.weight,
            None,
            self.stride,
            self.padding,
            self.dilation,
            self.groups)
        return out

class WeightModifier(torch.autograd.Function):
    def __init__(self, weight):
        self.weight = weight

    def forward(self, input):
        return input.clone()

    def backward(self, grad):
        # Some op to change self.weight.
        # In other words, change weights in the following conv2d
        # after loss.backward() and before optimizer.step()
        return grad.clone()

hidden = CustomConv()(input_tensor)  # backward not called 
hidden = CustomConv()(hidden)
loss = cal_loss(hidden, gt)
optimizer.zero_grad()
loss.backward()
optimizer.step()
</code></pre>
<p>The problem is the backward() of WeightModifier in the first CustomConv is not called (that in second CustomConv is called). I guess the reason is that Pytorch finds that input_tensor does not require gradient and layer WeightModifier does not have any parameters. Is the any method to force or “cheat” Pytorch to execute the backward()?</p>
<p>Thanks.</p>

albanD: <p>Hi,</p>
<p>The problem most likely comes from the fact that you use old style Functions. You can see <a href="https://pytorch.org/docs/stable/notes/extending.html" rel="nofollow noopener">in the doc</a> how to write new style ones.<br>
Do keep the <code>.clone()</code> in the forward though otherwise we detect that it’s the same Tensor and the <code>Function</code>'s backward won’t be called.</p>

hayden: <p>Hi albanD,<br>
Thanks. I still don’t make it work. As you said, I tried something like this:</p>
<pre><code class="lang-auto">class WeightModifier(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight):
        ctx.save_for_backward(weight)
        return input.clone()

    @staticmethod
    def backward(ctx, grad_output):
        weight, = ctx.saved_tensors
        weight.data.copy_(weight.org)
        return grad_output.clone()
</code></pre>
<p>which is called by <code>input = WeightModifier.apply(input, self.weight)</code>. However Pytorch requires backward() to return gradients for both input and weight, whereas I want it to return gradient for input only.</p>

hayden: <p>Sorry to bother you again. Here is another question. In <a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html" rel="nofollow noopener">this exampe</a>, input.clamp(min=0) is not recorded in gradient graph. Is it because it’s inside a forward function of autograd.Function? To generalize, any module/function inside a forward function of a Function does not contribute to gradient computation?</p>
<p>Thanks!</p>

albanD: <p>The whole point of using a <code>Function</code> is to not use the autograd.<br>
If you want to use the autograd, you can just use a regular python function.</p>
<p>If you don’t return gradients for one of the inputs, you can return <code>None</code> for it.</p>