Debugging autograd functions in pytorch

Hi! I wrote below function that I intend to use as the loss function.

class MyCriterion(torch.autograd.Function):
    def __init__(self):
        self.alpha = .0005
    
    def forward(self, input, target, epoch, isLabeled):
        loss = F.cross_entropy(input, target)

        self.save_for_backward(input, target, epoch, isLabeled, loss)
        print(self.saved_tensors) #returns ()
        if (isLabeled.data > 0).all():

            return Variable(loss.data * self.alpha * epoch.data)
        return loss

    def backward(self, grad_output):
       
        input, target, epoch, isLabeled, loss, = self.saved_tensors
        grad_input = loss.backward() 
        return grad_input
my_criterion = MyCriterion()
x = Variable(torch.randn(11, 10).type(torch.FloatTensor))
y = Variable(torch.range(1,6, .5).type(torch.LongTensor))

a = torch.from_numpy(np.array([0]))
b = torch.from_numpy(np.array([1]))
c = torch.from_numpy(np.array([10.0]))

print(x)
# print(torch.from_numpy(np.array([10])))
first_loss  = my_criterion.forward(x, y, Variable(c.float()),  Variable(a))
print(my_criterion.backward(first_loss))

second_loss = my_criterion.forward(x, y, Variable(c.float()),  Variable(b))
print(my_criterion.backward(second_loss))

When I do this, I have below error -

---> 18         input, target, epoch, isLabeled, loss, = self.saved_tensors
     19         grad_input = loss.backward()
     20         return grad_input

ValueError: not enough values to unpack (expected 5, got 0)

Is there something I am missing? How can I access saved tensors? Is there any other documentation than autograd that has more examples of autograd functions? Thanks a lot!

First thing is that you can’t return a Variable from forward - it expects tensors and will automatically wrap them in Variables and connect them up to the existing graph. You shouldn’t unpack and re-pack Variables in the middle of computation, because that will break continuity of history. You need to do something like that:

def my_criterion(input, target, epoch, isLabeled):
    if (isLabeled.data > 0).all():
        return loss * alpha * epoch
    return MyCriterion()(input, target, epoch, isLabeled)

Another thing. You should never call the forward method directly. You should instantiate the Function class and call it like you’d call a function. You can use the function I provided as a more convenient wrapper.

I’d recommend reading notes about extending autograd.

1 Like

It’s also unclear whether what you’re writing needs to be a Function. If you want the autograd library to automatically compute the backwards pass for your operation, and you can represent the operation as a combination of existing autograd-enabled functions (as it looks like you’ve done in forward), you should just define a Python function or a Module.

If you do need to write a new Function subclass, that means you aren’t able to represent the operation as a combination of existing functions with known derivatives, and you have to implement the backward pass yourself (you can’t just call .backward()). The computations inside the forward and backward methods of a Function subclass take place on Tensor objects, not Variables.

3 Likes

This is very useful, I will try writing it as a module, thanks!

A quick follow up question - since Module doesn’t have an explicit .backward() method, how do I exactly backprop on a loss function that is a module? Is it enough if I just use .train() instead?

Thanks!

all of the operators inside a Module's forward function have a backward defined, because the input is a Variable. So the backward for the module is automatically defined by autograd.