CODEWITHSUNDEEP
pythonmachine-learningneural-networkpytorchbackpropagation
Just_A_Doubt
Just_A_Doubt

Reputation: 13

Implementing Backprop for custom loss functions

I have a neural network Network that has a vector output. Instead of using a typical loss function, I would like to implement my own loss function that is a method in some class. This looks something like:

class whatever:
    def __init__(self, network, optimizer): 
         self.network = network
         self.optimizer = optimizer
    
    def cost_function(relevant_data):
        ...implementation of cost function with respect to output of network and relevant_data...

    def train(self, epochs, other_params):
        ...part I'm having trouble with...

The main thing I'm concerned with is about taking gradients. Since I'm taking my own custom loss function, do I need to implement my own gradient with respect to the cost function?

Once I do the math, I realize that if the cost is J, then the gradient of J is a fairly simple function in terms of the gradient of the final layer of the Network. I.e, it looks something like: Equation link.

If I used some traditional loss function like CrossEntropy, my backprocess would look like:

objective = nn.CrossEntropyLoss()
for epochs:
    optimizer.zero_grad()
    output = Network(input)
    loss = objective(output, data)
    loss.backward()
    optimizer.step()

But how do we do this in my case? My guess is something like:

for epochs:
     optimizer.zero_grad()
     output = Network(input)
     loss = cost_function(output, data)
     #And here is where the problem comes in
     loss.backward()
     optimizer.step()

loss.backward() as I understand it, takes the gradients of the loss function with respect to the parameters. But can I still invoke it while using my own loss function (presumably the program doesn't know what the gradient equation is). Do I have to implement another method/subroutine to find the gradients as well?

Which brings me to my other question: if I do want to implement gradient calculation for my loss function, I also need the gradient of the neural network parameters. How do I obtain those? Is there a function for that?

Upvotes: 1

Views: 2179

Answers (1)

GoodDeeds
GoodDeeds

Reputation: 8507

As long as all your steps starting from the input till the loss function involve differentiable operations on PyTorch's tensors, you need not do anything extra. PyTorch builds a computational graph that keeps track of each operation, its inputs, and gradients. So, calling loss.backward() on your custom loss would still propagate gradients back correctly through the graph. A Gentle Introduction to torch.autograd from the PyTorch tutorials may be a useful reference.

After the backward pass, if you need to directly access the gradients for further processing, you can do so using the .grad attribute (so t.grad for tensor t in the graph).

Finally, if you have a specific use case for finding the gradient of an arbitrary differentiable function implemented using PyTorch's tensors with respect to one of its inputs (e.g. gradient of the loss with respect to a particular weight in the network), you could use torch.autograd.grad.

Upvotes: 2

Related Questions