PyTorch参数化层无法更新求助：仅训练增量参数D失败

原代码的问题分析 #

第一种实现的问题 #

1. 参数维度不匹配：self.dw和self.db初始化为标量1.0，虽然PyTorch支持广播，但和frozen.weight（形状[10,100]）、frozen.bias（形状[10]）的维度逻辑不匹配，会导致后续计算的梯度传播异常。
2. forward中重复赋值Parameter切断梯度：每次前向传播都给self.fc.weight和self.fc.bias创建新的nn.Parameter，这会破坏原本的计算图关联，导致dw和db的梯度无法正确回传，自然不会被优化器更新。

第二种parametrize实现的问题 #

在forward方法中重复调用register_parametrization，每次前向传播都会给层注册新的参数化模块，导致参数集合混乱；且参数化注册应该在模型初始化阶段（__init__）完成，而非每次前向传播时执行。

正确实现方案 #

方案一：手动计算带增量的权重（无需parametrize） #

核心思路：不修改原有Linear层的权重属性，而是在forward中直接计算W = frozen.weight + dw、b = frozen.bias + db，手动执行线性变换，确保计算图完整。

import torch
import torch.nn as nn

# 初始化冻结权重层
frozen = nn.Linear(100, 10)
frozen.weight.requires_grad = False
frozen.bias.requires_grad = False

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 定义与冻结权重同形状的可训练增量参数，初始值可按需调整
        self.dw = nn.Parameter(torch.zeros_like(frozen.weight))
        self.db = nn.Parameter(torch.zeros_like(frozen.bias))

    def forward(self, x):
        # 计算当前使用的权重和偏置
        current_weight = frozen.weight + self.dw
        current_bias = frozen.bias + self.db
        # 手动执行线性变换，再经过激活函数
        out = torch.matmul(x, current_weight.T) + current_bias
        return torch.sigmoid(out)

# 训练流程
model = Net()
# CrossEntropyLoss需配合logits使用，若用sigmoid输出可改用BCEWithLogitsLoss
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# 调整输入为[batch_size, in_features]格式
x = torch.rand(1, 100)
# CrossEntropyLoss目标为类别索引（整数类型），而非one-hot向量
y = torch.tensor([9])

for _ in range(10):
    out = model(x)
    loss = criterion(out, y)
    print(f"Loss: {loss.item()}")
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    # 查看增量参数的更新情况
    print(f"dw均值: {model.dw.mean().item()}")

方案二：正确使用nn.utils.parametrize #

核心思路：在模型初始化阶段完成参数化注册，确保参数化模块仅被注册一次，增量参数被正确纳入模型的可训练参数集合。

import torch
import torch.nn as nn

# 初始化冻结权重层
frozen = nn.Linear(100, 10)
frozen.weight.requires_grad = False
frozen.bias.requires_grad = False

class Adder(nn.Module):
    def __init__(self, frozen_param):
        super().__init__()
        # 初始化与冻结参数同形状的可训练增量
        self.delta = nn.Parameter(torch.zeros_like(frozen_param))
        self.frozen_param = frozen_param

    def forward(self, x):
        # parametrize的forward接收原层参数，此处直接返回冻结参数+增量
        return self.frozen_param + self.delta

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(100, 10)
        # 在初始化阶段完成参数化注册，仅执行一次
        nn.utils.parametrize.register_parametrization(self.fc, "weight", Adder(frozen.weight))
        nn.utils.parametrize.register_parametrization(self.fc, "bias", Adder(frozen.bias))

    def forward(self, x):
        return torch.sigmoid(self.fc(x))

# 训练流程
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
x = torch.rand(1, 100)
y = torch.tensor([9])

for _ in range(10):
    out = model(x)
    loss = criterion(out, y)
    print(f"Loss: {loss.item()}")
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    # 查看权重增量的更新情况
    print(f"权重增量均值: {model.fc.parametrizations.weight[0].delta.mean().item()}")

额外注意事项 #

• 损失函数搭配：CrossEntropyLoss内置了log_softmax，若输出用sigmoid，需改用BCEWithLogitsLoss（二分类场景）或去掉sigmoid直接输出logits。
• 输入与目标格式：nn.Linear要求输入为[batch_size, in_features]；CrossEntropyLoss的目标需为类别索引（整数张量），而非one-hot向量。

内容的提问来源于stack exchange，提问作者malioboro