with torch.no_grad() 和 叶子节点in place操作

with torch.no_grad() 和 叶子节点in place操作

requires_grad为true的叶子节点不能进行in place操作

例如下面的代码

w = torch.Tensor([10.]) # w 是个 leaf tensor
w.requires_grad = True    # 将 requires_grad 设置为 True
y = w**2
w.normal_()               # 在执行这句话就会报错
# 报错信息为
#  RuntimeError: a leaf Variable that requires grad has been used in an in-place operation.

叶子节点不能执行in-place(原地)操作，因为在进行前向传播的时候得到的是叶子结点的地址，再进行反向传播的时候这个地址不变才不会报错。所以pytorch设置为不能改变。

另一种情况是requires_grad为true非叶子节点，在求梯度时需要被用到，也不可以进行in place操作。

import torch
x = torch.FloatTensor([[1., 2.]])
w1 = torch.FloatTensor([[2.], [1.]])
w2 = torch.FloatTensor([3.])
w1.requires_grad = True
w2.requires_grad = True

d = torch.matmul(x, w1)
f = torch.matmul(d, w2)
d[:] = 1 # 因为这句报错

f.backward() # 报错指向这里 RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation

但这里如果修改的是f也不会报错，只有中间节点不能修改。
参考https://zhuanlan.zhihu.com/p/38475183

动手学深度学习中的with torch.no_grad()

在该模块下，所有计算得出的tensor的requires_grad都自动设置为False。

即使一个tensor（命名为x）的requires_grad = True，在with torch.no_grad计算，由x得到的新tensor（命名为w-标量）requires_grad也为False，且grad_fn也为None,即不会对w求导。

这在验证和测试环节非常有用，可以节省内存，不必要开辟计算图。但是动手学深度学习中有不少训练环节也出现了这样的代码，分别解析作用

3.2.6节

def sgd(params, lr, batch_size):  #@save
    """小批量随机梯度下降"""
    with torch.no_grad():
        for param in params:
            param -= lr * param.grad / batch_size
            param.grad.zero_()

这里面params对应的是线性回归里的[w, b]

一方面，如果去掉torch.no_grad()，这个代码直接运行不了了。因为param是叶子节点，而-=(__iadd__)是in place操作，所以会报错（上面第一种情况）。

另一方面，从目的上讲，我们就是希望对param做原位操作。如果我们去掉torch.no_grad(), 改成param = param - lr * param.grad / batch_size，也可以运行。但是param就成了计算图中新的节点，在调用sgd的外层循环里，进入下一次循环时，不再是叶子节点，也不会有梯度信息，整个计算图都乱掉了。

其实也可以不用no grad，像参考网站medium那篇博客一样，在反向传播之前加上retain_grad参数

import torch

X = torch.tensor([1, 2, 3, 4], dtype=torch.float32)
Y = torch.tensor([2, 4, 6, 8], dtype=torch.float32)

w = torch.tensor(0.0, dtype=torch.float32, requires_grad=True)

learning_rate = 0.01
n_iters = 100

for epoch in range(n_iters):
    # predict = forward pass
    y_pred = w * X
    # loss
    l = (y_pred - Y).pow(2).mean()
    # calculate gradients = backward pass
    w.retain_grad()
    l.backward()
    
    w = w - learning_rate * w.grad
    # zero the gradients after updating
    w.grad = None
    print(f'epoch {epoch+1}: w = {w.item():.3f}, loss = {l.item():.8f}')

也可以运行，就是内存开销太大了

3.2.7节

for epoch in range(num_epochs):
    for X, y in data_iter(batch_size, features, labels):
        l = loss(net(X, w, b), y)  # X和y的小批量损失
        # 因为l形状是(batch_size,1)，而不是一个标量。l中的所有元素被加到一起，
        # 并以此计算关于[w,b]的梯度
        l.sum().backward()
        sgd([w, b], lr, batch_size)  # 使用参数的梯度更新参数
    with torch.no_grad():
        train_l = loss(net(features, w, b), labels)
        print(f'epoch {epoch + 1}, loss {float(train_l.mean()):f}')

这个不难理解，就是评测的时候为了减少内存开销。

参考网站

https://medium.com/@mrityu.jha/understanding-the-grad-of-autograd-fc8d266fd6cf

https://discuss.pytorch.org/t/leaf-variable-was-used-in-an-inplace-operation/308/2

https://stackoverflow.com/questions/71322901/questions-on-using-grad-within-torch-no-grad