如何修复PyTorch中因原地操作导致的梯度计算错误

问题背景 #

实现了论文中“单力作用的一维系统”（case 1）的代码，目标是用神经网络近似物理系统的力。反向传播过程中出现梯度计算错误，通过torch.autograd.set_detect_anomaly(True)定位到错误关联adsmodel类_forward方法的ans=(v_c-v_abs)*self.force[v_f]+(v_abs-v_f)*self.force[v_c]行。

错误信息：

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [10]], which is output 0 of SubBackward0, is at version 81; expected version 80 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!

问题根源 #

错误本质并非_forward方法的计算行，而是**forward方法中对输入tensor的原地修改操作**：v-=(self.g-(self._forward(v)/self.m))*self.delta_t。原地操作会破坏PyTorch的计算图追踪链路，导致梯度计算时变量版本不匹配。

修复方案 #

关键修改点 #

1. 替换原地操作：避免直接修改输入tensor，创建新变量追踪计算图
2. 同步设备一致性：确保索引tensor与参数在同一设备（CPU/GPU）
3. 移除冗余参数：删除不必要的retain_graph=True，减少内存占用

修复后完整代码 #

import torch
from torch import nn
import numpy as np

class adsmodel(nn.Module):
    def __init__(self, delta_t, t_s, t_f, m, g, max_v, batch_size, c):
        super().__init__()
        self.delta_t = delta_t
        self.m = m
        self.g = g
        self.t_f = t_f
        self.t_s = t_s
        self.batch_size = batch_size
        self.L = max_v + 1
        self.c = c
        self.force = nn.Parameter(torch.FloatTensor(np.random.uniform(low=10, high=20, size=(self.L,))))

    def _forward(self, v):
        v_abs = torch.abs(v)
        # 确保索引tensor与参数设备一致
        v_c = torch.ceil(v_abs).type(torch.long).to(v.device)
        v_f = torch.floor(v_abs).type(torch.long).to(v.device)
        ans = (v_c - v_abs) * self.force[v_f] + (v_abs - v_f) * self.force[v_c]
        return ans
    
    def forward(self, v):
        # 克隆输入tensor，避免原地修改破坏计算图
        current_v = v.clone()
        steps = int((self.t_f - self.t_s) / self.delta_t)
        for _ in range(steps):
            force_term = self._forward(current_v) / self.m
            update = (self.g - force_term) * self.delta_t
            # 非原地操作，创建新节点追踪梯度
            current_v = current_v - update
        return current_v
    
    def loss(self, x, y):
        result = self.forward(x)
        regular = 0
        for i in range(self.L - 1):
            regular += (self.force[i+1] - self.force[i]) ** 2
        regular = regular * self.c + self.force[0] ** 2
        loss = torch.mean(torch.abs(result - y)) + regular
        return loss

def train(epoch, model, train_x, train_y, lr):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    for i in range(epoch):
        print(f"Epoch {i+1}/{epoch}")
        loss = model.loss(train_x, train_y)
        optimizer.zero_grad()
        # 移除retain_graph=True，默认自动释放计算图
        loss.backward()
        optimizer.step()
        print(f"Loss: {loss.item():.4f}")
        
torch.autograd.set_detect_anomaly(True)
        
model = adsmodel(0.1, 0, 4, 1, 9.8, 500, 10, 0.03)
train_x = torch.FloatTensor(np.random.uniform(-100, -10, (10,)))
train_y = torch.FloatTensor(np.random.uniform(-200, -50, (10,)))
print(f"Initial Loss: {model.loss(train_x, train_y).item():.4f}")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
train_x, train_y = train_x.to(device), train_y.to(device)

train(100, model, train_x, train_y, 0.01)

修复说明 #

• 原地操作替换：通过current_v = current_v - update创建新的tensor节点，确保计算图完整追踪梯度流向
• 设备同步：索引tensor通过.to(v.device)与输入tensor设备对齐，避免GPU环境下的索引错误
• 内存优化：移除retain_graph=True，每次迭代后自动释放计算图，减少不必要的内存占用

内容的提问来源于stack exchange，提问作者8기장민균