PyTorch设置requires_grad=False在GPU运行时无法冻结网络参数问题

问题根源 #

你的问题本质是优化器的参数列表初始化逻辑、Adam优化器的动量特性，和动态修改requires_grad的时序不匹配，具体可以拆解为两点：

1. 优化器的参数列表是初始化时一次性确定的：你在代码最开始初始化Adam优化器时，所有参数的requires_grad都是True，所以Sub_Net的参数已经被加入了优化器的更新列表，后续你再动态修改requires_grad为False，也不会把这些参数从优化器的更新列表中移除。
2. Adam的动量更新逻辑在梯度为0时依然会生效：你每次循环先执行generator_step，此时Sub_Net的参数是可训练的，反向传播后会生成grad属性，同时Adam优化器会为这些参数累计动量（一阶矩m_t）和二阶矩v_t。进入discrimination_step后，你把Sub_Net参数的requires_grad设为False，但参数的grad属性依然存在，你调用net.zero_grad()只会把grad的值设为0，而不会把grad清空为None。此时反向传播不会更新这些参数的梯度，梯度保持为0，Adam会按照动量规则继续更新参数：m_t = beta1 * m_{t-1} + (1-beta1)*0，v_t = beta2 * v_{t-1} + (1-beta2)*0，依然会产生非零的更新量，这就是你GPU输出中逐渐减小的差值的来源。
   CPU上输出全为0是因为CPU端的浮点计算精度、PyTorch算子实现的微小差异，导致更新量小到8位小数下被截断为0，不是真的没有更新。

解决方案 #

推荐两种可靠的修复方式：

方式1：使用两个独立优化器（最稳妥） #

分别为两个step初始化对应的优化器，从根源上避免discrimination_step更新Sub_Net的参数：

# 初始化两个优化器
optimizer_generator = optim.Adam(net.parameters(), lr=0.1)
# discrimination_step的优化器只包含Sub_Net以外的参数
optimizer_discrimination = optim.Adam([p for n,p in net.named_parameters() if not n.startswith('Sub_Net.')], lr=0.1)

然后在对应的step里使用对应的优化器即可。

方式2：动态清空冻结参数的梯度 #

如果你要保留单优化器的逻辑，在discrimination_step中设置requires_grad=False后，额外把Sub_Net参数的grad设为None，让优化器直接跳过这些参数的更新：

def discrimination_step(net, optimizer, criterion, input, target):
    for param in net.Sub_Net.parameters():
        param.requires_grad = False
        # 额外把grad设为None，避免优化器用历史动量更新
        param.grad = None
    # 其余代码保持不变

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