如何在PyTorch中为每个参数设置独立的可学习学习率？

我来帮你搞定这个问题——你的思路方向是对的，但之前的方案在参数注册和更新逻辑上有两个关键漏洞，导致学习率参数没法被正确训练。下面我会一步步拆解问题，给出可行的实现方案。

先说说你之前方案的问题 #

1. 可学习LR未被注册为模型参数：你把学习率参数放在了普通的OrderedDict里，PyTorch的nn.Module只会把直接作为类属性的nn.Parameter（或nn.ParameterDict/nn.ParameterList）纳入可训练参数集合，字典里的参数不会被model.parameters()识别，自然没法被optimizer优化。
2. 参数更新逻辑冲突：你手动用LR更新模型参数后又调用了optimizer.step()，这会让optimizer用默认的args.lr再更新一次模型参数，覆盖了你手动的操作；同时因为LR参数没被注册，optimizer也不会去优化它们。

正确的实现方案 #

1. 正确注册可学习学习率参数 #

我们用nn.ParameterDict来存储每个模型参数对应的可学习LR，它是PyTorch专门设计的参数容器，会自动把里面的nn.Parameter纳入模型的可训练参数列表：

import torch
import torch.nn as nn
from collections import OrderedDict

class Learner(nn.Module):
    def __init__(self, net, in_channels, num_classes):
        super().__init__()
        self.learner = net(in_channels, num_classes)
        # 用ParameterDict存储每个参数对应的可学习LR
        self.param_lrs = nn.ParameterDict()
        for name, param in self.learner.named_parameters():
            # 初始化LR：这里用1e-4的固定值，你也可以根据需求用正态/均匀初始化
            lr_param = nn.Parameter(torch.full_like(param, 1e-4), requires_grad=True)
            self.param_lrs[name] = lr_param

    def forward(self, X):
        return self.learner(X)

2. 自定义双阶段优化逻辑 #

我们需要分开处理模型参数和学习率参数的更新：

• 模型参数：手动用对应的可学习LR进行更新
• 学习率参数：用单独的optimizer来优化，让它们在训练中随损失自动调整

下面是完整的训练流程示例（假设你的基础模型是一个简单CNN）：

# 定义基础模型示例
class SimpleNet(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.layer1 = nn.Conv2d(in_channels, 32, kernel_size=3)
        self.layer2 = nn.Linear(32*26*26, num_classes)  # 适配28x28输入图像

    def forward(self, x):
        x = torch.relu(self.layer1(x))
        x = x.flatten(1)
        x = self.layer2(x)
        return x

# 初始化带可学习LR的模型
model = Learner(SimpleNet, in_channels=1, num_classes=10)
# 只给学习率参数定义optimizer，模型参数我们手动更新
lr_optimizer = torch.optim.Adam(model.param_lrs.parameters(), lr=0.001)

# 训练循环
for epoch in range(10):
    # 假设你已经有数据加载器dataloader
    for X, y in dataloader:
        # 清零所有梯度
        lr_optimizer.zero_grad()
        model.zero_grad()

        # 前向传播计算损失
        outputs = model(X)
        loss = nn.CrossEntropyLoss()(outputs, y)

        # 反向传播：同时计算模型参数和学习率参数的梯度
        loss.backward()

        # 手动用可学习LR更新模型参数（关闭梯度追踪）
        with torch.no_grad():
            for name, param in model.learner.named_parameters():
                if param.grad is not None:
                    # 用对应参数的可学习LR完成更新
                    param.data -= model.param_lrs[name].data * param.grad.data

        # 更新学习率参数本身
        lr_optimizer.step()

        # 打印训练状态
        print(f"Epoch {epoch+1}, Batch Loss: {loss.item():.4f}")

方案核心说明 #

• 参数注册：nn.ParameterDict确保学习率参数被PyTorch识别为模型的可训练参数，这样model.param_lrs.parameters()就能正确获取这些LR参数交给optimizer优化。
• 梯度关联：当你用model.param_lrs[name] * param.grad更新模型参数时，这个操作会被PyTorch的自动微分系统追踪，从而让学习率参数获得对应的梯度——也就是说，损失会引导LR参数朝着减少损失的方向调整。
• 无冲突更新：我们用单独的optimizer只优化LR参数，模型参数则手动更新，避免了默认optimizer的干扰。

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