Conv2d层单权重指定置零及不可更新冻结操作实现方法咨询

PyTorch的requires_grad是张量级别的属性，无法直接给张量的某个子集单独设置该属性，也不能直接对作为Leaf节点的模型参数做原位赋值修改，要实现细粒度的权重置零+冻结，可使用如下两种成熟方案：

方案1：梯度裁剪钩子+预初始化置零（推荐，代码侵入性最低） #

• 第一步：先通过.data属性修改参数的实际数值，避开计算图校验完成指定位置置零：

import torch
import torch.nn as nn

model = nn.Conv2d(3,3,kernel_size=5, stride=1)
# 操作.data属性不会触发计算图记录，可正常完成指定位置赋值
model.weight.data[0,2] = 0

• 第二步：给权重参数注册反向传播钩子，梯度回传后自动把指定位置的梯度置零，等价于冻结效果，后续优化器更新时这些位置的权重不会发生变化：

def weight_grad_hook(grad):
    # 按需修改需要冻结的位置的梯度为0即可
    grad[0,2] = 0
    return grad

model.weight.register_hook(weight_grad_hook)

后续按原有逻辑正常训练即可，你指定的位置会始终保持0值且不会被更新。如果需要批量冻结多个位置，提前生成和权重形状一致的布尔掩码，反向时执行grad[mask] = 0即可。

方案2：拆分权重参数（适合冻结规则固定的场景） #

如果冻结位置是长期固定的，也可以把卷积核拆分为「需要冻结」和「需要更新」两个独立的参数，仅给需要更新的部分设置requires_grad=True，前向传播时拼接为完整的卷积核再做卷积计算即可，该方案需要自行重写卷积层的前向逻辑。

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