如何通过全连接层生成1x1卷积参数且不破坏PyTorch autograd计算图

梯度无法回流的根因 #

你当前的写法破坏Autograd计算图有两个核心原因：

• with torch.no_grad() 上下文管理器会主动阻断张量的梯度传播，包裹参数赋值操作后，全连接层的输出到卷积层参数的路径完全没有梯度记录
• 直接替换nn.Conv2d层的weight属性属于模块参数替换操作，新的nn.Parameter不会保留和之前全连接层输出的计算关联，梯度自然无法回传到上游的全连接网络

修复方案 #

1x1卷积的本质是通道维度的线性变换，你不需要实例化nn.Conv2d模块再替换参数，直接使用PyTorch的函数式卷积接口torch.nn.functional.conv2d传入全连接层输出的权重即可，全程张量运算都会被Autograd自动追踪，不会断梯度。

同时你原代码中存在两处明显错误：一是__init__方法没有调用父类构造函数，二是定义的input_layers、conv_layers在forward中被写成了input_layer、conv_layer，都需要同步修正。另外1x1卷积权重形状为[输出通道数, 输入通道数, 1, 1]，你原全连接层的最终输出维度不符合要求，也需要对应调整。

修改后可运行代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F

class Network(nn.Module):
    def __init__(self, len_input, num_kernels):
        super().__init__()
        self.input_layers = nn.Sequential(
            nn.Linear(len_input, num_kernels * 2),
            nn.ReLU(),
            nn.Linear(num_kernels * 2, num_kernels * 3), # 输出维度匹配1x1卷积权重总参数量：num_kernels*3*1*1
            nn.ReLU()
        )

        self.conv_layers = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(in_channels=num_kernels, out_channels=3, kernel_size=1)
        )

    def forward(self, x1, img):
        # 全连接层输出动态生成卷积权重
        x = self.input_layers(x1.float())
        # 权重reshape为[out_channels, in_channels, 1, 1] 对应1x1卷积的权重格式
        conv_weight = x.reshape(-1, 3, 1, 1)
        # 用函数式接口做卷积运算，全程保留梯度路径
        conv_out = F.conv2d(img, conv_weight, bias=None, stride=1, padding=0)
        generated = self.conv_layers(conv_out)
        return generated

原全连接层参数直方图 #

内容的提问来源于stack exchange，提问作者Furkan Küçük