Manual backprob of gradient in last layer to loss

solmazst: <p>Hi,</p>
<p>I have a model called ‘NET’ which is bunch convolutional layers.<br>
as my loss function I am using soft-dtw which has a work-around to make DTW differentiable. the gradient of soft-dtw is a matrix with the size of my input data (x).</p>
<p>x: input data<br>
Net: CNN<br>
loss: stdw<br>
derivative of loss w.r.t. y = G (matrix)<br>
y = Net(x)  --&gt; output data</p>
<p>normally when we calculate loss and do loss.backward(), the gradient of the Loss on last layer w.r.t model params is calculated and back-propagated through the model parameters. however here, my loss function is not a pytorch loss library , it’s a self-define function (soft-dtw). To calculate the gradient of loss w.r.t Net.parameters, I use the chain rule:</p>
<p>derivative of Loss w.r.t  y    * derivative of y w.r.t  model’s parameters<br>
here if I define a Loss class that is inheriting from autograd, and do Loss.backward, is there a way to pass G to loss.backward() so that it uses G as the gradient of last layer and multiply it by the gradient of y w.r.t. model params?</p>
<pre><code class="lang-auto">y = Net(x)
loss = SDTWLoss(x,y)
 # gradient of loss w.r.t y is  Loss_grad = G  ,a matrix where G.shape = x.shape
loss.backward()
optimizer.step()


# right now in my loss function I have implemented the code as the following:
class SDTWLoss():
    
    def __init__(self, y_pred, y):
        self.y_pred = y_pred
        self.y = y
        
    def forward(self):
        
        _dtw = SoftDTW(self.y_pred, self.y ,device)
        dtw_value, dtw_grad, dtw_E_mat= _dtw.compute()
        self.dtw_grad = dtw_grad
        dtw_loss = torch.mean(dtw_value)
        
        return dtw_loss

    
    def backward(self):
    
        batch_size, _, n = self.y_pred.shape
        G = jacobian_product_with_expected_sdtw_grad(self.y_pred, self.y, self.dtw_grad,device)
     
        param_grad = []
        for param in Net.parameters():
             param_grad.append(param)
        
        for k in range(batch_size):
            for i in range(n):
                Net.zero_grad()
                self.y_pred[k,0,i].backward(retain_graph=True)
                for j,param in enumerate(Net.parameters()):
                     param_grad[j] =param_grad[j]+ (G[k,0,i] * param.grad)


        for j,param in enumerate(Net.parameters()):
            param.grad = param_grad[j]
</code></pre>
<p>this code is extremely slow!<br>
I think the fastest way is to have the option of passing G as an input to loss.backward:<br>
‘’’    loss.backward(G) ‘’’’<br>
Any suggestion is appreciated!</p>

cosmozhang1988: <p>Have you tried this?<br>
<a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html" class="onebox" target="_blank" rel="nofollow noopener">https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html</a></p>

solmazst: <p>my understanding is that this example is implementing Relu activation function.  I need to implement Loss function for specific case where the last layer’s derivative is calculated manually and is back propagated through the network using chain rule.</p>

cosmozhang1988: <p>Yeah. This is a <code>relu</code> example, but you can extend this to your own problem. The returned value in <code>backward</code> are the derivatives of the input of the <code>forward</code>, which is the output of the second to the last layer in your case. Your loss is implemented in the <code>forward</code>, or constructed by using the output of the <code>forward</code>. You just need to write your last layer’s derivative manually in the <code>backward</code>.</p>

solmazst: <p>thanks for clarifying. I am still not completely there.<br>
so let’s say I have the following:</p>
<p>‘’’<br>
class SDTWLoss(torch.autograd.Function):</p>
<pre><code>@staticmethod
def forward( ctx, y_pred, y):
   
    _dtw = SoftDTW(y_pred, y ,device)
    dtw_value, dtw_grad, dtw_E_mat= _dtw.compute() # calculate loss value for y_prediction and y_actual
    G = jacobian_product_with_expected_sdtw_grad(y_pred, y, dtw_grad,device) # calculate derivative of loss w.r.t. y_pred (last layer of my graph) 
    ctx.save_for_backward(G) #save derviative of Loss w.r.t y for backward
    dtw_loss = torch.mean(dtw_value)

    return dtw_loss


@staticmethod
def backward( ctx, grad_output ):

    G = ctx.saved_tensors
    grad_input = grad_output.clone() # is this the derivative of the layer before the last layer w.r.t model's parameters???
    # if so, I should multiply G by grad_input here, right?
    grad_input = G*grad_input

    return grad_input
</code></pre>
<p>‘’’<br>
Does this code make sense?</p>

cosmozhang1988: <p>The logic is correct here.</p>

solmazst: <p>unfortunately it’s not working yet!<br>
if grad_output is the gradient w.r.t y_predicted, I am calculating it manually (matrix G). what should I multiply it to , to get the gradient w.r.t. model’s params?</p>

cosmozhang1988: <p>You need to do <code>return grad_input, None</code> in your <code>backward</code> as you do not need grads for true <code>y</code>.</p>