【附代码案例】深入理解 PyTorch 张量：叶子张量与非叶子张量

在 PyTorch 中，张量是构建神经网络模型的基本元素。了解张量的属性和行为对于深入理解模型的运行机制至关重要。本文将介绍 PyTorch 中的两种重要张量类型：叶子张量和非叶子张量，并探讨它们在反向传播过程中的行为差异。

叶子张量与非叶子张量的区别

1. 叶子张量是由用户直接创建的张量，而非叶子张量是通过对叶子张量进行操作得到的张量。可以通过 .is_leaf 属性来判断一个张量是否是叶子节点。

2. 叶子张量是需要求梯度的张量，因此它们会保存计算图的结构以便进行反向传播。而非叶子张量一般是通过张量的加减乘除、函数的调用等操作得到的，它们不会保存计算图的结构，因此不会自动求梯度。

3. 默认情况下，对于 requires_grad=True 的张量，默认情况下，它们是叶子张量。

非叶子张量的梯度累积

对于非叶子张量，每次调用 loss.backward() 后，梯度并不会清零，而是会累积到对应张量的 .grad 属性中。这意味着梯度会在反向传播过程中持续累积，直到显式清零。

优化器的梯度清零方法

优化器的 optimizer.zero_grad_() 方法可以将优化器中所有参数张量的梯度清零，包括叶子张量和非叶子张量。这样做的目的是为了防止梯度的累积，确保每一次反向传播都是基于当前 batch 的梯度计算而不会受之前 batch 的影响。

requires_grad 属性的作用

requires_grad 是一个布尔值属性，用于指示张量是否需要计算梯度。如果 requires_grad 为 True，则 PyTorch 会在张量上的操作中跟踪梯度信息，允许通过调用 .backward() 方法自动计算梯度。默认情况下，张量的 requires_grad 属性为 False。

查看梯度的方法

在执行反向传播之后，可以通过访问张量的 .grad 属性来查看梯度。在反向传播之前，这些张量的梯度值是不存在的，因此打印出来的是 None。如果希望在非叶子节点张量上累积梯度，需要在计算前调用 .retain_grad() 方法。

通过深入理解叶子张量与非叶子张量的区别以及它们在反向传播过程中的行为，可以更好地掌握 PyTorch 的工作机制，并有效地调试和优化神经网络模型。

代码示例

下面是一个简单的示例，演示了如何使用 PyTorch 创建叶子张量和非叶子张量，并观察它们在反向传播过程中的行为：

import torch

# 创建叶子张量
leaf_tensor = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# 创建非叶子张量
non_leaf_tensor = leaf_tensor * 2

# 求非叶子张量的平方和作为损失函数
loss = torch.sum(non_leaf_tensor ** 2)

# 打印非叶子张量是否是叶子节点
print("non_leaf_tensor is leaf:", non_leaf_tensor.is_leaf)

# 调用反向传播计算梯度
loss.backward()

# 查看叶子张量的梯度
print("Gradient of leaf_tensor:", leaf_tensor.grad)

# 查看非叶子张量的梯度
print("Gradient of non_leaf_tensor:", non_leaf_tensor.grad)

# 再次调用反向传播计算梯度，梯度会累积
loss.backward()

# 查看叶子张量的梯度
print("Gradient of leaf_tensor after second backward:", leaf_tensor.grad)

# 查看非叶子张量的梯度
print("Gradient of non_leaf_tensor after second backward:", non_leaf_tensor.grad)

在这个示例中，我们首先创建了一个叶子张量 leaf_tensor，然后通过对其进行操作得到了一个非叶子张量 non_leaf_tensor。我们使用 non_leaf_tensor 的平方和作为损失函数，然后调用反向传播计算梯度。可以观察到，虽然 non_leaf_tensor 是由 leaf_tensor 操作得到的，但它的梯度仍然会被计算并存储在 .grad 属性中。