20天狂宴Pytorch-Day7

发表于 2025-09-02 更新于 2025-09-03

Pytorch 的动态计算图.

动态计算图#

Pytorch 的计算图由节点和边组成, 节点表示张量或者 Function, 边表示张量和 Function 之间的依赖关系.

Pytorch 的计算图是动态图:

计算图的正向传播是立即执行的, 无需等待完整的计算图创建完毕, 每条语句都会在计算图中动态添加节点和边, 并立即执行正向传播得到计算结果.
计算图在反向传播后立即销毁. 如果在程序中使用了 backward 方法执行了反向传播, 或者利用torch.autograd.grad方法计算了梯度, 那么创建的计算图会被立即销毁, 下次调用需要重新创建.

import torch 
w = torch.tensor([[3.0,1.0]],requires_grad=True)
b = torch.tensor([[3.0]],requires_grad=True)
X = torch.randn(10,2)
Y = torch.randn(10,1)
Y_hat = X @ w.t() + b
# Y_hat定义后其正向传播被立即执行，与其后面的loss创建语句无关

loss = torch.mean(torch.pow(Y_hat-Y,2))
print(loss.data)
print(Y_hat.data)

loss.backward()
loss.backward() # 再次反向传播将报错

计算图中的 Function#

计算图中的另外一种节点是 Function, 即 Pytorch 中各种对张量操作的函数, 但其同时包括正向计算的逻辑和反向传播的逻辑.

可以通过继承torch.autograd.Function来创建这种支持反向传播的 Function.

class MyReLU(torch.autograd.Function):
   
    # 正向传播逻辑
    # 可以用ctx存储一些值供反向传播使用
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    # 反向传播逻辑
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

import torch 
w = torch.tensor([[3.0,1.0]],requires_grad=True)
b = torch.tensor([[3.0]],requires_grad=True)
X = torch.tensor([[-1.0,-1.0],[1.0,1.0]])
Y = torch.tensor([[2.0,3.0]])

relu = MyReLU.apply
Y_hat = relu(X@w.t() + b)
loss = torch.mean(torch.pow(Y_hat-Y,2))
loss.backward()
print(w.grad)
print(b.grad)
# Y_hat的梯度函数即MyReLU.backward
print(Y_hat.grad_fn)

计算图与反向传播#

import torch 

x = torch.tensor(3.0,requires_grad=True)
y1 = x + 1
y2 = 2*x
loss = (y1-y2)**2

loss.backward()

调用loss.backward()后, 依次发生:

从 loss 标量出发, loss 自身的 grad 梯度赋值为 1
loss 根据其自身梯度以及关联的 backward 方法, 计算出其对应的自变量 (y1 和 y2) 的梯度, 将该值赋值到 y1.grad 和 y2.grad
y1 和 y2 根据其自身梯度以及关联的 backward 方法, 计算出其对应的自变量 (x) 的梯度, x.grad 将收到的多个梯度值累加

因此张量的 grad 梯度不会自动清零, 需要手动清零.

叶子节点和非叶子节点#

import torch 

x = torch.tensor(3.0,requires_grad=True)
y1 = x + 1
y2 = 2*x
loss = (y1-y2)**2

loss.backward()
print("loss.grad:", loss.grad)
print("y1.grad:", y1.grad)
print("y2.grad:", y2.grad)
print(x.grad)

执行上述代码发现, loss.grad 并非期望的 1, 而是 None, 类似地 y1.grad 和 y2.grad 也是 None. 原因是它们不是叶子节点张量, 在反向传播过程中, 只有is_leaf=True的节点, 导数结果才会被保留.

叶子节点张量需要满足两个条件:

是由用户直接创建的张量, 而非由某个 Function 通过计算得到的张量
requires_grad 属性必须为 True

所有依赖于叶子节点张量的张量, 其 requires_grad 属性必定为 True, 但其梯度值只在计算过程中被用到, 不会最终存储到 grad 属性中.

如果需要保留中间计算结果的梯度到 grad 属性中, 可以使用 retain_grad 方法; 如果仅仅是为了调试代码查看梯度值, 可以利用 register_hook 打印日志. 利用 retain_grad 可以保留非叶子节点的梯度值, 利用 register_hook 可以查看非叶子节点的梯度值.

import torch 

# 正向传播
x = torch.tensor(3.0,requires_grad=True)
y1 = x + 1
y2 = 2*x
loss = (y1-y2)**2

# 非叶子节点梯度显示控制
y1.register_hook(lambda grad: print('y1 grad: ', grad))
y2.register_hook(lambda grad: print('y2 grad: ', grad))
loss.retain_grad()

# 反向传播
loss.backward()
print("loss.grad:", loss.grad)
print("x.grad:", x.grad)