Error in the backward of custom loss function

username: bill
cooked: <p>Hi, I’m new in the pytorch.<br>
I have a question about the custom loss function.<br>
The code is following.<br>
I use numpy to clone the MSE_loss as MSE_SCORE.<br>
Input is 1x200x200 images, and batch size is 128.<br>
The output “mse” of the MSE_SCORE is a float value and covered as tensor.</p>
<p>However, when I tried to apply the backward of custom loss function, then got the error “RuntimeError: Function MSE_SCOREBackward returned an invalid gradient at index 0 - got [] but expected shape compatible with [128, 1, 200, 200]”</p>
<pre><code class="lang-python"> class MSE_SCORE(Function):
     @staticmethod
     def forward(ctx, input, label):
         with torch.no_grad():
             numpy_input = input.detach().numpy()
             numpy_label = label.detach().numpy()
             offset_height = int(height / 4)
             offset_width = int(width / 4)
             target_height = int(height / 2)
             target_width = int(width / 2)
             crop_input = numpy_input[:, :, offset_height:offset_height + target_height, offset_width:offset_width + target_width]
             crop_label = numpy_label[:, :, offset_height:offset_height + target_height, offset_width:offset_width + target_width]
             diff = (crop_input - crop_label)**2
             mean = np.mean(diff, axis=(1,2,3))
             mse = np.sum(mean)
 
         return torch.tensor(mse)
 
     @staticmethod
     def backward(ctx, grad_output):
         return grad_output, None
 
 def criterion_mse(input, label):
     return MSE_SCORE.apply(input, label)
 
 for epoch in range(num_epochs):
     for step, (img, label) in enumerate(dataloader):
         img = Variable(img).to(device=device)
         label = Variable(label).to(device=device)
         output = model(img)
         mse_loss = criterion_mse(output.cpu(), label.cpu()).to(device=device)
         optimizer.zero_grad()
         loss_seq = [mse_loss]
         grad_seq = [torch.tensor(1.0).to(device=device) for _ in range(len(loss_seq))]
         torch.autograd.backward(loss_seq, grad_seq)</code></pre>

username: albanD
cooked: <p>Hi,</p>
<p>The problem here is that your implementation of the backward is not correct.</p>
<ul>
<li>You don’t need to use torch.no_grad() is a custom Function</li>
<li>Why not run the exact same code (without the conversion to numpy) in your criterion_mse() function so that it will be autodiffed and you get backward for free?</li>
<li>If you don’t do that, you need to implement the backward yourself. And in particular, the error you see here is because the gradient your return for <code>input</code> don’t have the same size as `input.</li>
</ul>
<p>PS: I updated your post to make the code more readable using triple backticks.</p>

username: bill
cooked: <p>Hi albanD,<br>
Thank a lot for reply.</p>
<p>Because I want to do more complex work for this custom loss function, I have to use the conversion to numpy.<br>
I’m not sure how grad_output works, do you have any example?<br>
The following code modified by myself. Is it correct?</p>
<pre><code class="lang-auto">class MSE_SCORE(Function):
    @staticmethod
    def forward(ctx, input, label):
        numpy_input = input.detach().numpy()
        numpy_label = label.detach().numpy()
        offset_height = int(height / 4)
        offset_width = int(width / 4)
        target_height = int(height / 2)
        target_width = int(width / 2)
        crop_input = numpy_input[:, :, offset_height:offset_height + target_height, offset_width:offset_width + target_width]
        crop_label = numpy_label[:, :, offset_height:offset_height + target_height, offset_width:offset_width + target_width]
        diff = (crop_input - crop_label)**2
        mean = np.mean(diff, axis=(1,2,3))
        ctx.size = len(mean)
        mse = np.sum(mean)
        ctx.save_for_backward(input, label) --- modified
        return torch.tensor(mse)

    @staticmethod
    def backward(ctx, grad_output):
        grad_output = grad_output.detach()
        input, label = ctx.saved_tensors --- modified
        return grad_output * input, None --- modified
</code></pre>

username: albanD
cooked: <p>The <code>backward()</code> function should compute one step of the chain rule of your function <code>y = f(x)</code><br>
<code>grad_output</code> is the gradient flowing from the lower layers <code>dl/dy</code>. And you should return <code>dl/dx</code>.<br>
The way to compute it is by doing: <code>dl/dx = dl/dy * dy/dx</code>.</p>

username: bill
cooked: <p>Hi <a class="mention" href="/u/alband">@albanD</a></p>
<p>I don’t know about “grad_output” means. Accounding to “dl/dx = dl/dy * dy/dx”, does it means dy / dx?</p>
<p>With my poor understanding:<br>
l = sum(mean( (label - input)**2 ))<br>
dl / dx = -2 * sum(mean( (label - input) )) * dy / dx<br>
Is it correct?</p>
<pre><code class="lang-auto">class MSE_SCORE(Function):
    @staticmethod
    def forward(ctx, input, label):
        numpy_input = input.detach().numpy()
        numpy_label = label.detach().numpy()
        diff = (numpy_label - numpy_input)**2
        mean = np.mean(diff, axis=(1,2,3))
        mse = np.sum(mean)
        ctx.save_for_backward(input, label)

        return torch.tensor(mse)

    @staticmethod
    def backward(ctx, grad_output):
        grad_output = grad_output.detach()
        input, label = ctx.saved_tensors
        numpy_input = input.detach().numpy()
        numpy_label = label.detach().numpy()
        grad = -2 * np.sum(np.mean((numpy_label - numpy_input), axis=(1,2,3))) * np.ones(input.shape)
        return grad_output * grad, None
</code></pre>

username: albanD
cooked: <p>Hi,</p>
<p><code>grad_output </code> is the gradient wrt to the output of the forward, so in this case <code>dl/dy</code>.</p>
<p>And yes the formula looks good.</p>
<p>In the implementation,</p>
<ul>
<li>for the forward, even though you can go through numpy, I don’t think you should here: all these ops just work on pytorch.</li>
<li>For the backward, you should write it in a differentiable manner in case you try to access grad of grad in there. If you don’t, use <code>from torch.autograd.function import once_differentiable</code> as another decorator for the backward.</li>
<li>Same point for the backward that I don’t see any reason to use numpy, even if you make it <code>once_differentiable</code>.</li>
</ul>