如何摆脱PyTorch.autograd中的Variable API？

Question

I am forwarding, and backpropping tensor data X through two simple nn.Module PyTorch models instances, model1 and model2 .

I can't get this process to work without usage of the depreciated Variable API .

So this works just fine:

    y1 = model1(X)
    v = Variable(y1.data, requires_grad=training)         # Its all about this line!
    y2 = model2(v)
    criterion = nn.NLLLoss()
    loss = criterion(y2, y)
    loss.backward()
    y1.backward(v.grad)
    self.step()

But this will throw an error:

    y1 = model1(X)
    y2 = model2(y1)
    criterion = nn.NLLLoss()
    loss = criterion(y2, y)
    loss.backward()
    y1.backward(y1.grad) # it breaks here
    self.step()
>>> RuntimeError: grad can be implicitly created only for scalar outputs

I just can't seem to find a relevant difference between v in the first implementation, and y1 in the second. In both cases requires_grad is set to True . The only thing I could find was that y1.grad_fn=<ThnnConv2DBackward> and v.grad_fn=<ThnnConv2DBackward>

What am I missing here? What (tensor attributes?) do I not know about, and if Variable is depreciated, what other implementation would work?

Answer 1

[UPDATED] You are not correctly passing the y1.grad into y1.backward in the second example. After the first backward all the intermediate gradient will be destroyed, you need a special hook to extract that gradients. And in your case you are passing the None value. Here is small example to reproduce your case:

Code:

import torch
import torch.nn as nn


torch.manual_seed(42)


class Model1(nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x.pow(3)


class Model2(nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x / 2


model1 = Model1()
model2 = Model2()
criterion = nn.MSELoss()

X = torch.randn(1, 5, requires_grad=True)
y = torch.randn(1, 5)

y1 = model1(X)
y2 = model2(y1)

loss = criterion(y2, y)
# We are going to backprop 2 times, so we need to 
# retain_graph=True while first backward
loss.backward(retain_graph=True)

try:
    y1.backward(y1.grad)
except RuntimeError as err:
    print(err)
    print('y1.grad: ', y1.grad)

Output:

grad can be implicitly created only for scalar outputs
y1.grad:  None

So you need to extract them correctly:

Code:

def extract(V):
    """Gradient extractor.
    """
    def hook(grad):
        V.grad = grad
    return hook


model1 = Model1()
model2 = Model2()
criterion = nn.MSELoss()

X = torch.randn(1, 5, requires_grad=True)
y = torch.randn(1, 5)

y1 = model1(X)
y2 = model2(y1)

loss = criterion(y2, y)
y1.register_hook(extract(y1))
loss.backward(retain_graph=True)

print('y1.grad', y1.grad)

y1.backward(y1.grad)

Output:

y1.grad:  tensor([[-0.1763, -0.2114, -0.0266, -0.3293,  0.0534]])

Answer 2

After some investigation I came to the following two solutions. The solution provided elsewhere in this thread retained the computation graph manually, without an option the free them, thus running fine initially, but causing OOM errors later on.

The first solution is to tie the models together using the built in torch.nn.Sequential as such:

model = torch.nn.Sequential(Model1(), Model2())

it's as easy as that. It looks clean and behaves exactly like an ordinary model would.

The alternative is to simply tie them together manually:

        model1 = Model1()
        model2 = Model2()
        y1 = model1(X)
        y2 = model2(y1)
        loss = criterion(y2, y)
        loss.backward()

My fear that this would only backpropagate model2 turned out to be unsubstantiated, since model1 is also stored in the computation graph that is back propagated over. This implementation enabled inceased transparancy of the interface between the two models, compared to the previous implementation.