PyTorch训练中no grad错误的诊断与修复

在pytorch训练过程中遇到`runtimeerror: element 0 of tensors does not require grad`错误，通常是由于计算损失的张量不具备梯度追踪能力所致。这可能是因为在计算图中的关键点执行了不可微分操作（如`argmax`），或者不当使用了`torch.no_grad()`。解决此问题的核心在于确保损失函数直接作用于模型输出的logits，并避免在梯度反向传播路径上引入会中断计算图的操作。

理解RuntimeError: element 0 of tensors does not require grad

当PyTorch抛出RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn错误时，意味着在执行loss.backward()时，PyTorch的自动微分系统（Autograd）无法找到一个需要计算梯度的张量，或者这个张量没有关联的梯度函数（grad_fn）。这通常发生在以下几种情况：

1. 张量没有设置requires_grad=True： 默认情况下，用户创建的张量（如输入数据、标签）不追踪梯度。只有模型参数和经过可微分操作产生的张量才会自动设置requires_grad=True。
2. 计算图被中断： 在计算损失的路径上，执行了某些不可微分的操作（如argmax、将张量转换为Python数字、使用.detach()方法等），或者将张量转换为不追踪梯度的类型（如int）。
3. 误用torch.no_grad()或torch.inference_mode()： 这些上下文管理器会暂时禁用梯度计算。如果训练代码的核心部分（尤其是损失计算）被错误地包裹在其中，就会导致此错误。

诊断训练循环中的问题

分析提供的训练循环代码：

for epoch in range(Epochs):
    model.train()

    train_logits = model(X_train)
    # 问题所在：argmax 操作中断了计算图
    train_preds_probs = torch.softmax(train_logits,dim=1).argmax(dim=1).type(torch.float32)
    loss = loss_fn(train_preds_probs,y_train) # 损失函数接收的是硬预测标签
    train_accu = accuracy(y_train,train_preds_probs)
    print(train_preds_probs)
    optimiser.zero_grad()

    loss.backward() # 此时 loss 的输入 train_preds_probs 已不追踪梯度

    optimiser.step()

    # ... 省略评估部分 ...

核心问题在于这一行： train_preds_probs = torch.softmax(train_logits,dim=1).argmax(dim=1).type(torch.float32)

1. torch.softmax(train_logits, dim=1)：这一步仍然保留了梯度信息。
2. .argmax(dim=1)：这是一个离散操作。它返回的是索引，而不是连续的值。argmax操作是不可微分的，它会从计算图中移除其输入（train_logits）的梯度追踪能力。
3. .type(torch.float32)：即使将结果转换回浮点数，也无法恢复已被argmax中断的梯度追踪。

因此，当loss = loss_fn(train_preds_probs,y_train)计算损失时，train_preds_probs已经是一个不具备requires_grad=True属性的张量，并且不关联任何grad_fn。随后调用loss.backward()时，系统发现无法对loss进行反向传播，因为它的输入不追踪梯度，从而抛出错误。

解决方案

解决此问题的关键在于确保损失函数直接作用于模型输出的原始logits，而不是经过argmax处理后的硬预测标签。对于分类任务，常用的交叉熵损失函数（如torch.nn.CrossEntropyLoss）通常期望模型的原始logits作为输入，并自动在内部执行softmax和负对数似然计算。

修正后的训练循环示例：

Ghiblio

专业AI吉卜力风格转换平台，将生活照变身吉卜力风格照

157

查看详情

import torch
import torch.nn as nn
import torch.optim as optim

# 假设的模型和数据
class SimpleModel(nn.Module):
    def __init__(self, input_dim, num_classes):
        super().__init__()
        self.linear = nn.Linear(input_dim, num_classes)

    def forward(self, x):
        return self.linear(x)

# 示例数据
input_dim = 10
num_classes = 3
batch_size = 32
X_train = torch.randn(batch_size, input_dim, requires_grad=True) # 模拟输入数据
y_train = torch.randint(0, num_classes, (batch_size,)) # 模拟标签
X_test = torch.randn(batch_size, input_dim)
y_test = torch.randint(0, num_classes, (batch_size,))

model = SimpleModel(input_dim, num_classes)
loss_fn = nn.CrossEntropyLoss() # 使用 nn.CrossEntropyLoss，它直接接受 logits 和整数标签
optimiser = optim.Adam(model.parameters(), lr=0.01)

# 辅助函数：计算准确率
def accuracy(y_true, y_pred_logits):
    y_pred_labels = torch.softmax(y_pred_logits, dim=1).argmax(dim=1)
    return (y_pred_labels == y_true).float().mean()

"""Training"""
Epochs = 100

for epoch in range(Epochs):
    model.train()

    train_logits = model(X_train)
    # 正确做法：损失函数直接作用于 logits 和真实标签
    loss = loss_fn(train_logits, y_train)

    # 准确率计算可以继续使用 argmax，因为它不需要梯度反向传播
    train_accu = accuracy(y_train, train_logits)

    optimiser.zero_grad()
    loss.backward() # 现在 loss 的计算路径上都是可微分操作，可以正常反向传播
    optimiser.step()

    # 评估阶段
    model.eval()
    with torch.inference_mode(): # 在评估阶段使用 torch.inference_mode() 是正确的
        test_logits = model(X_test)
        test_loss = loss_fn(test_logits, y_test) # 评估损失同样作用于 logits
        test_acc = accuracy(y_test, test_logits)

    if epoch % 10 == 0:
        print(f'Epoch:{epoch} | Train loss: {loss.item():.4f} | Taining acc:{train_accu:.4f} | Test Loss: {test_loss.item():.4f} | Test accu: {test_acc:.4f}')

关键改动说明：

• loss = loss_fn(train_logits, y_train)：将损失函数loss_fn的输入从train_preds_probs（经过argmax处理）改回train_logits（模型的原始输出）。nn.CrossEntropyLoss内部会处理softmax和对数似然计算，因此不需要在外部手动执行softmax。
• train_accu = accuracy(y_train, train_logits)：准确率的计算可以继续使用argmax，因为准确率本身不需要参与梯度反向传播。它只是一个指标。

最佳实践与注意事项

1. 损失函数选择：
   • 对于多分类任务，torch.nn.CrossEntropyLoss是首选，它接受模型的原始logits和整数类型的真实标签。
   • 对于二分类任务，torch.nn.BCEWithLogitsLoss是首选，它同样接受原始logits和浮点类型的真实标签（0或1）。
   • 避免手动在损失函数外部执行softmax后再传递给CrossEntropyLoss，这可能导致数值不稳定。
2. 梯度追踪的边界：
   • 明确哪些操作会中断梯度追踪：.detach()、.item()、.numpy()、int()、argmax()等。
   • 只有在确定不需要对某个张量进行反向传播时，才使用.detach()。
   • 在打印张量值或将其用于非梯度计算时，可以使用.item()获取Python数值，但这也会中断梯度追踪。
3. torch.no_grad()的正确使用：
   • 仅在评估模型、进行推理或计算不需要梯度的指标时使用with torch.no_grad():或with torch.inference_mode():。
   • 确保训练循环中涉及梯度计算的部分（模型前向传播、损失计算、loss.backward()）不在这些上下文管理器内部。
4. 数据类型：
   • 确保模型输出和损失函数期望的标签数据类型匹配。例如，nn.CrossEntropyLoss期望logits为float类型，标签为long类型。
   • 在需要时使用.type()或.to()进行类型转换，但要注意这可能会创建新的张量，如果需要梯度追踪，需确保新张量也具备requires_grad=True。

总结

RuntimeError: element 0 of tensors does not require grad错误的核心在于梯度反向传播路径上的张量失去了梯度追踪能力。在PyTorch训练中，解决此问题的关键是确保：

1. 损失函数的输入是具备梯度追踪能力的张量，通常是模型的原始输出（logits）。
2. 避免在损失计算路径上执行不可微分操作（如argmax），这些操作会中断计算图。
3. 正确使用torch.no_grad()或torch.inference_mode()，仅在不需要梯度计算的场景下使用它们。

遵循这些原则，可以有效地避免此类错误，并构建健壮的PyTorch训练流程。

以上就是PyTorch训练中no grad错误的诊断与修复的详细内容，更多请关注php中文网其它相关文章！