ResNet替换卷积层参数后，如何正确获取自定义参数A、B的梯度？

你的代码存在三个关键错误，导致反向传播后A和B的梯度始终为None：

1. B被重赋值，丢失叶子节点属性：执行B = B.view(-1, 9)后，原nn.Parameter对象被覆盖成view操作后的张量，不再是叶子节点，后续的requires_grad_和retain_grad操作对原Parameter无效。
2. base被包装为新Parameter，断开计算图：将base转为nn.Parameter赋值给卷积层时，这个新Parameter是独立的叶子节点，和A、B的计算图完全断开，反向传播不会回溯到A和B。
3. A和B未加入模型参数列表：PyTorch仅对注册到模型中的参数计算梯度，你的A和B是独立于model的张量，反向传播时不会被处理。

以下是修正后的代码，核心思路是保留A和B的叶子节点属性，将它们注册到模型中，让卷积层权重动态依赖A和B的计算结果：

import torch
import torch.nn as nn
from torchvision import models

# 初始化ResNet18
model = models.resnet18()

# 将A、B注册为模型的参数属性
model.A = nn.Parameter(torch.zeros(size=(64 * 64, 32 * 32)))
model.B = nn.Parameter(torch.zeros(size=(32, 32, 3, 3)))

# 定义前向钩子：每次前向传播前用A、B更新卷积层权重
def update_conv_weight(module, input):
    # 对B做维度变换，不覆盖原Parameter对象
    B_flat = model.B.view(-1, 9)
    # 计算卷积权重的基础张量
    base = model.A @ B_flat
    base = base.view(64, 64, 3, 3)
    # 直接修改卷积层权重的data，保留原Parameter结构
    module.weight.data = base.data

# 给目标卷积层注册前向钩子
hook_handle = model.layer1[0].conv1.register_forward_pre_hook(update_conv_weight)

# 测试数据
img = {
    "img": torch.randn(size=(1, 3, 28, 28)),
    "label": torch.ones(1).long()
}

# 前向传播+计算损失（模拟你的loss逻辑）
output = model(img["img"])
loss = nn.CrossEntropyLoss()(output, img["label"])

# 反向传播
loss.backward()

# 现在可以正常获取A、B的梯度
print("A.grad:", model.A.grad)
print("B.grad:", model.B.grad)

# 用完钩子后移除
hook_handle.remove()

• 将A、B注册为模型属性：通过model.A = nn.Parameter(...)让PyTorch把它们纳入模型参数管理，反向传播时会自动计算梯度。
• 用前向钩子动态更新权重：在每次前向传播前，用A、B计算卷积权重并赋值给module.weight.data，既保留原卷积层的Parameter结构，又让权重依赖A、B的计算图。
• 避免重赋值Parameter：对B做维度变换时，直接生成新张量但不覆盖原model.B，保留原叶子节点的引用。

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