How does Pytorch no_grad function for a = a - b and a -= b type of operation?

Viewed 100
import torch

def model(x, W, b):
    return x@W + b


def mse(t1, t2):
    diff = t1 - t2
    return torch.sum(diff * diff) / diff.numel()

inputs = torch.rand(2, 3, requires_grad=True)
targets = torch.rand( 2,2, requires_grad=True)

W = torch.rand(3, 2, requires_grad=True)
b = torch.rand(2, requires_grad=True)

pred = model(inputs, W, b)

loss = mse(pred, targets)
loss.backward()

print(W.grad)
print(b.grad)
with torch.no_grad():
    W -= W.grad * 1e-5
    b -= b.grad * 1e-5
print(W.grad)
print(b.grad)

For the above example, the output of the last 2 print statements is the same as that of the first two print statements.

But, for the code snippet below, the 2 last print statements give result as None. I can't understand why it is so?

print(W.grad)
print(b.grad)
with torch.no_grad():
    W = W - W.grad * 1e-5
    b = b -  b.grad * 1e-5
print(W.grad)
print(b.grad)
1 Answers

Keep in mind in both scenarios you are under the torch.no_grad context manager which by effect disables gradient computation.

On one hand, you are performing an in-place operation on your tensor which means their underlying data gets modified without changing the reference two that tensor storage in memory, moreover its metadata remains unchanged, that is W and b are and remain tensors which require gradient (as defined in the very first assignments with requires_grad=True).

On the other, you are performing out-of-place operations which means variables W and b both get assigned brand new tensors. Indeed, out-of-place assignments create copies. Therefore the W and b are no longer the ones defined prior to the assignment but different ones. Not only their values are different, but the tensors' metadata itself has changed. Finally, the reason why you have None is that tensors defined under this context manager will not have a requires_grad=True set by definition of the context.

Related