PyTorch自定义含可训练/固定参数模型训练报错求助

错误原因 #

你当前的代码在__init__方法中预先计算了matrix张量，这个张量与可训练参数w绑定了计算图。第一次反向传播（cost.backward()）后，PyTorch会自动释放计算图的中间张量以节省内存。第二次迭代时，再次使用这个已经失去计算图的matrix进行前向传播，后续反向传播就会报错，因为无法追溯到参数w的梯度路径。

修复方案 #

将矩阵的计算逻辑移到forward方法中，这样每次前向传播都会重新构建与当前参数状态绑定的计算图，每次反向传播只会销毁当前批次的计算图，不会影响下一次迭代。

修改后的完整代码：

import torch

class CustomModel(torch.nn.Module):
    def __init__(self, w0):
        super(CustomModel, self).__init__()
        self.w = torch.nn.Parameter(data=torch.tensor([w0], dtype=torch.float32, requires_grad=True))
        # 提前定义固定矩阵，避免每次forward重复创建
        self.fixed_mat = torch.tensor([[0, 1], [-1, 0]], dtype=torch.float32)
        self.eye_mat = torch.eye(2, dtype=torch.float32)
    
    def forward(self, x):
        # 每次forward动态计算带可训练参数的矩阵
        matrix_trainable = self.w * self.fixed_mat
        matrix = matrix_trainable - self.eye_mat
        return matrix.matmul(x)

def loss(pred, y):
    return torch.mean((pred - y)**2)

my_model = CustomModel(w0=0.01)
optimizer = torch.optim.Adam(lr=0.01, params=my_model.parameters())

device = torch.device("cpu")
x = torch.ones(2).to(device)
y = torch.tensor([2., 0.], dtype=torch.float32).to(device)

for k in range(10):
    optimizer.zero_grad()
    pred = my_model(x)
    cost = loss(pred, y)
    cost.backward()
    optimizer.step()
    print(f"Epoch {k+1}, Loss: {cost.item():.6f}, w: {my_model.w.item():.6f}")

额外说明 #

• 固定矩阵（比如fixed_mat、eye_mat）可以在__init__中提前创建，避免每次forward重复初始化，提升效率。
• 如果确实需要在__init__中保留矩阵定义，也可以通过register_buffer方法注册固定张量，但可训练部分的计算必须放在forward中，确保每次迭代都有有效的计算图。

内容的提问来源于stack exchange，提问作者Matías Bilkis