4、动态计算图
4、动态计算图
4.1、动态计算图简介
计算图(Computational Graph)是一个有向无环图(DAG),计算图由节点和边组成
- 节点:张量或者Function
- 边:张量和 Function 之间的依赖关系
计算图是动态图,动态图有两层含义:
- 计算图的 正向传播是立即执行 的。
无需等待完整的计算图创建完毕,每条语句都会在计算图中动态添加节点和边,并立即执行正向传播得到计算结果 - 计算图在反向传播后立即销毁。
下次调用需要重新构建计算图。
如果在程序中使用了 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)
计算图在反向传播后立即销毁
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))
# 计算图在反向传播后立即销毁,如果需要保留计算图, 需要设置retain_graph = True
loss.backward() #loss.backward(retain_graph = True)
# loss.backward() # 如果再次执行反向传播将报错
保存计算图多次反向传播
import torch
# 创建需要梯度的张量
x = torch.tensor(1.0, requires_grad=True)
y = torch.tensor(2.0, requires_grad=True)
# 第一次计算和反向传播
z = x * y
w = z * y
w.backward(retain_graph=True) # 保留计算图以便后续反向传播
# 查看第一次反向传播后的梯度
print(f"x.grad after first backward: {x.grad}") # 输出:tensor(4.0)
print(f"y.grad after first backward: {y.grad}") # 输出:tensor(2.0)
# 第二次计算和反向传播
v = x * y
u = v * y
u.backward()
# 查看第二次反向传播后的梯度
print(f"x.grad after second backward: {x.grad}") # 输出:tensor(8.0)
print(f"y.grad after second backward: {y.grad}") # 输出:tensor(6.0)
4.2、计算图中的 Function
计算图中的另外一种节点是 Function, 实际上就是 Pytorch 中各种对张量操作的函数
这些 Function 同时包括正向计算逻辑和反向传播的逻辑
可以通过继承 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 # relu 现在也可以具有正向传播和反向传播功能
Y_hat = relu(X@w.t() + b)
loss = torch.mean(torch.pow(Y_hat-Y,2))
loss.backward()
print(w.grad) # tensor([[4.5000, 4.5000]])
print(b.grad) # tensor([[4.5000]])
# Y_hat的梯度函数即是我们自己所定义的 MyReLU.backward
print(Y_hat.grad_fn)
4.3、计算图与反向传播
简单地理解一下反向传播的原理和过程
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 自己的 grad 梯度赋值为 1,即对自身的梯度为 1
loss 根据其自身梯度以及关联的 backward 方法,计算出其对应的自变量即 y1 和 y2 的梯度,将该值赋值到 y1.grad 和 y2.grad
y2 和 y1 根据其自身梯度以及关联的 backward 方法, 分别计算出其对应的自变量 x 的梯度,x.grad 将其收到的多个梯度值累加
(注意,1,2,3步骤的求梯度顺序和对多个梯度值的累加规则恰好是求导链式法则的程序表述)
正因为求导链式法则衍生的梯度累加规则,张量的 grad 梯度不会自动清零,在需要的时候需要手动置零
4.4、叶子节点和非叶子节点
举例,下面图中的叶子结点与非叶子结点:
- 叶子结点:w,x,b
- 非叶子结点:y,z

在反向传播过程中,只有 is_leaf=True 的 叶子节点,需要求导的张量的导数结果才会被最后保留下来。
叶子节点张量需要满足两个条件。
- 叶子节点张量是 由用户直接创建 的张量,而非由某个 Function 通过计算得到的张量
- 叶子节点张量的 requires_grad 属性必须为 True
Pytorch 设计这样的规则主要是为了节约内存或者显存空间,因为几乎所有的时候,用户只会关心他自己直接创建的张量的梯度。
所有依赖于叶子节点张量的张量, 其 requires_grad 属性必定是 True 的,但其梯度值只在计算过程中被用到,不会最终存储到 grad 属性中。
如果需要保留中间计算结果的梯度到 grad 属性中,可以使用 retain_grad 方法。 如果仅仅是为了调试代码查看梯度值,可以利用 register_hook 打印日志。
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) # loss.grad: None
print("y1.grad:", y1.grad) # y1.grad: None
print("y2.grad:", y2.grad) # y2.grad: None
print(x.grad) # tensor(4.)
print(x.is_leaf) # True
print(y1.is_leaf) # False
print(y2.is_leaf) # False
print(loss.is_leaf) # False
利用 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)
'''
y2 grad: tensor(4.)
y1 grad: tensor(-4.)
loss.grad: tensor(1.)
x.grad: tensor(4.)
'''
4.5、计算图在 TensorBoard 中的可视化
from torch import nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.w = nn.Parameter(torch.randn(2,1))
self.b = nn.Parameter(torch.zeros(1,1))
def forward(self, x):
y = x@self.w + self.b
return y
net = Net()
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('./data/tensorboard')
writer.add_graph(net,input_to_model = torch.rand(10,2))
writer.close()
%load_ext tensorboard
#%tensorboard --logdir ./data/tensorboard
from tensorboard import notebook
notebook.list()
# 在tensorboard中查看模型
notebook.start("--logdir ./data/tensorboard")