PyTorch中如何为特定神经网络参数单独缩放梯度？

有两种常用方法可以实现你要的需求，分别适用于不同场景：

方法一：反向传播后手动修改梯度值 #

这种方法直接操作参数的梯度张量，灵活性更高，适合需要动态调整缩放比例的场景。

步骤说明 #

1. 正常执行前向传播计算损失
2. 反向传播生成所有参数的梯度
3. 定位到目标参数（参数C），将其梯度乘以指定的缩放因子
4. 调用优化器的step()完成参数更新

代码示例 #

import torch
import torch.nn as nn
import torch.optim as optim

# 定义包含参数A、B、C的神经网络
class MyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc_a = nn.Linear(10, 20)  # 参数A所在层
        self.fc_b = nn.Linear(20, 30)  # 参数B所在层
        self.fc_c = nn.Linear(30, 5)   # 参数C所在层（包含weight和bias）

net = MyNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)

# 训练循环示例
for inputs, labels in your_dataloader:
    optimizer.zero_grad()
    
    # 前向传播
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    
    # 反向传播生成梯度
    loss.backward()
    
    # 缩放参数C的梯度（这里设置为2倍，可替换为1/3）
    scale_factor = 2.0
    for param in net.fc_c.parameters():
        if param.grad is not None:
            param.grad *= scale_factor
    
    # 更新所有参数
    optimizer.step()

如果参数C是单个独立参数（而非某层的所有参数），只需定位到该参数直接修改梯度即可：

# 假设net中有单个参数c_param
net.c_param.grad *= scale_factor

方法二：通过优化器参数分组调整学习率 #

因为参数的更新量公式为 更新量 = 学习率 × 梯度，所以给目标参数设置缩放后的学习率，等价于直接缩放梯度的更新效果。这种方法代码更简洁，适合固定缩放比例的场景。

代码示例 #

# 给不同参数分组设置学习率，参数C的学习率为基础学习率的2倍
optimizer = optim.SGD([
    {'params': net.fc_a.parameters()},  # 使用基础学习率0.01
    {'params': net.fc_b.parameters()},  # 使用基础学习率0.01
    {'params': net.fc_c.parameters(), 'lr': 0.01 * 2}  # 参数C使用2倍学习率
], lr=0.01)

# 训练循环无需额外修改梯度，正常执行即可
for inputs, labels in your_dataloader:
    optimizer.zero_grad()
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

两种方法的对比 #

• 方法一：直接操作梯度，支持动态调整缩放因子（比如根据训练轮次、损失值变化调整），但需要额外的梯度修改步骤
• 方法二：通过学习率间接实现，代码更简洁，但缩放比例固定时更适用

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