自定义层

中文版 pytorch
42

原先的代码存在一个重要问题：在 forward() 中使用 .data 会导致权重无法正常参与自动求导，模型无法通过正常的反向传播学习这些参数。下面是两种写法的对比

import torch
from torch import nn
from torch.nn import functional as F

class MyLinear(nn.Module):
    def __init__(self, use_data):
        super().__init__()

        self.weight = nn.Parameter(torch.tensor([[1.0], [2.0]]))
        self.bias = nn.Parameter(torch.tensor([0.0]))
        self.use_data = use_data

    def forward(self, X):
        if self.use_data:
            # 错误写法
            linear = X @ self.weight.data + self.bias.data
        else:
            # 正确写法
            linear = X @ self.weight + self.bias

        return F.relu(linear)


X = torch.tensor([[1.0, 1.0]])

for use_data in [True, False]:
    print(f"\n使用 .data：{use_data}")

    net = MyLinear(use_data)
    Y = net(X)
    loss = Y.sum()

    print("loss:", loss.item())
    print("loss.requires_grad:", loss.requires_grad)

    try:
        loss.backward()
    except RuntimeError as e:
        print("反向传播失败:", e)

    print("weight.grad:", net.weight.grad)
    print("bias.grad:", net.bias.grad)