叶子张量和非叶子张量

叶子张量（Leaf Tensor）和非叶子张量（Non-leaf Tensor）是 PyTorch 计算图中的两个核心概念，核心区别体现在创建方式、在计算图中的角色以及梯度存储行为上，具体差异如下：

1. 创建方式不同

• 叶子张量：通过直接初始化创建（如 torch.tensor()、torch.zeros()、torch.randn() 等），不是任何张量运算的结果。
  例：x = torch.tensor([1.0], requires_grad=True)（直接创建，叶子张量）。

• 非叶子张量：通过张量运算生成（如加减乘除、矩阵乘法等），是计算图中中间步骤的结果。
  例：y = x + 2（由 x 运算生成，非叶子张量，即使 y.requires_grad=True）。

2. 在计算图中的角色不同

• 叶子张量：是计算图的“起点”，通常对应模型的输入数据（如 x）或需要优化的参数（如权重 w、偏置 b）。它们没有“父节点”（不依赖其他张量）。

• 非叶子张量：是计算图的“中间节点”，依赖于其他张量（叶子张量或其他非叶子张量）。例如模型的中间层输出、损失函数的中间计算结果等。

3. 梯度存储行为不同（核心差异）

• 叶子张量：调用 backward() 后，其梯度会被自动存储在 .grad 属性中（前提是 requires_grad=True），供后续优化器更新（如 optimizer.step()）。

• 非叶子张量：默认情况下，梯度会被自动丢弃（.grad 为 None），因为它们不是需要优化的参数，保留梯度会浪费内存。若需保留，需手动设置 retain_grad()。

4. requires_grad 的默认行为

• 叶子张量：requires_grad 默认为 False，若需求导需显式设为 True（如模型参数）。
• 非叶子张量：requires_grad 由依赖的张量决定——只要有一个依赖的张量 requires_grad=True，则非叶子张量的 requires_grad 自动为 True（无需手动设置）。

示例对比

import torch

# 叶子张量：直接创建，requires_grad=True
x = torch.tensor([2.0], requires_grad=True)  
w = torch.tensor([3.0], requires_grad=True)  

# 非叶子张量：由运算生成（x*w 和 x*w + 1 都是运算）
y = x * w       # 非叶子张量，requires_grad=True（因 x 和 w 可求导）
z = y + 1       # 非叶子张量，requires_grad=True

# 反向传播：计算 z 对所有叶子张量的梯度
z.backward()

# 叶子张量的 .grad 有值
print(x.grad)  # tensor([3.])（dz/dx = w = 3）
print(w.grad)  # tensor([2.])（dz/dw = x = 2）

# 非叶子张量的 .grad 默认为 None
print(y.grad)  # None
print(z.grad)  # None

# 若需保留非叶子张量的梯度，需手动调用 retain_grad()
y.retain_grad()
z.backward(retain_graph=True)  # 需保留计算图，否则第二次求导会报错
print(y.grad)  # tensor([1.])（dz/dy = 1）

总结

• 叶子张量：计算图的起点，需手动创建，梯度会被保存，用于参数优化。
• 非叶子张量：计算图的中间结果，由运算生成，梯度默认不保存，仅用于前向传播和反向传播的链式计算。

理解这一区别，能帮助你更清晰地掌控 PyTorch 的自动求导流程，避免梯度无法获取或内存浪费等问题。