PyTorch多模态模型线性层参数不更新问题求助

问题背景 #

• 现有配备5×5卷积核的CNN，已针对特定扩散系数优化，可输出温度图
• 尝试构建简单前馈网络，输入不同扩散系数以生成对应卷积核，进而输出正确温度图
• 核心问题：前馈网络的两个线性层参数在训练过程中完全无更新，损失值始终保持一致

问题复现代码 #

from load_model import *
import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

n_train = 15

input_folder = f"data/ALLMAPS/iteration_no0"
output_folder = f"data/ALLMAPS/iteration_no10"
dataset = HeatDiffusion_multi_alpha_random(input_folder, output_folder)

n_samples = len(dataset)
n_test = n_samples-n_train
train_set, test_set = torch.utils.data.random_split(dataset, [n_train,n_test])

class FrozenConv2d(nn.Conv2d):
    def __init__(self):
        super().__init__(in_channels=1, out_channels=1, kernel_size=(5, 5), padding=2, padding_mode='replicate',
                         bias=False)
        self.weight.requires_grad = False  # freeze the convolution kernel
        # self.bias.requires_grad = False

    def forward(self, x):
        out = nn.functional.conv2d(x, self.weight, bias=None, padding=2)  # , self.bias)
        return out
conv_layer = FrozenConv2d()

class Smart(nn.Module):
    def __init__(self):
        super(Smart, self).__init__()
        self.l1 = nn.Linear(1,5)
        self.l2 = nn.Linear(5,25)
        self.act = nn.ReLU()

    def forward(self, x):
        alpha = x[1]
        alpha = alpha.view(-1,1)
        pre_kernel = self.act((self.l1(alpha)))
        kernel = self.l2(pre_kernel).view(1,1,5,5)

        conv_layer.weight = nn.Parameter(kernel) #requires_grad is set to true otherwise the backward is not working

        image = x[0].view(-1, 1, 100, 100)
        out = self.act((conv_layer(image)))
        out = out.view(100,100)

        return out

model = Smart().to(device)
model.train()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.003) #here model.parameters() only contains l1 and l2 which I want to optimize

for inputs, true_outputs in train_set:
    optimizer.zero_grad()
    inputs = [inp.to(device) for inp in inputs]
    true_outputs = true_outputs.to(device)

    # forward
    pred_outputs = model(inputs)
    loss = criterion(pred_outputs, true_outputs)

    # backwards
    loss.backward()
    optimizer.step()

print('loss', loss.item())
layers1 = [x.data for x in model.parameters()]

for inputs, true_outputs in train_set:
    optimizer.zero_grad()
    inputs = [inp.to(device) for inp in inputs]
    true_outputs = true_outputs.to(device)

    # forward
    pred_outputs = model(inputs)
    loss = criterion(pred_outputs, true_outputs)

    # backwards
    loss.backward()
    optimizer.step()

print('loss', loss.item()) #the loss is exactly the same
layers2 = [x.data for x in model.parameters()] # the parameters are exactly the same

问题症状 #

• 训练前后线性层l1、l2的参数完全一致
• 两轮训练后损失值无任何变化
• 确认线性层参数requires_grad=True，但梯度无法反向传播更新参数
• 未使用批处理，原因是原生Conv2d不支持单批内每个样本使用不同卷积核

问题根源与解决方案 #

核心问题：计算图断裂 #

代码中conv_layer.weight = nn.Parameter(kernel)的操作，是将前馈网络生成的kernel赋值给外部定义的Conv层实例，而该实例不属于Smart模型的子模块。PyTorch的自动微分机制无法追踪这种跨模块的参数赋值，导致梯度无法流回前馈网络的线性层，最终参数无法更新。

修复方案：将卷积操作整合到模型内部 #

移除外部的FrozenConv2d类，直接在Smart模型的forward方法中使用nn.functional.conv2d完成卷积计算，确保计算图完整连接线性层与卷积输出，梯度可正常反向传播。

修改后的完整代码：

from load_model import *
import torch
import torch.nn as nn
import torch.nn.functional as F

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

n_train = 15

input_folder = f"data/ALLMAPS/iteration_no0"
output_folder = f"data/ALLMAPS/iteration_no10"
dataset = HeatDiffusion_multi_alpha_random(input_folder, output_folder)

n_samples = len(dataset)
n_test = n_samples-n_train
train_set, test_set = torch.utils.data.random_split(dataset, [n_train,n_test])

class Smart(nn.Module):
    def __init__(self):
        super(Smart, self).__init__()
        self.l1 = nn.Linear(1,5)
        self.l2 = nn.Linear(5,25)
        self.act = nn.ReLU()

    def forward(self, x):
        alpha = x[1]
        alpha = alpha.view(-1,1)
        pre_kernel = self.act(self.l1(alpha))
        # 生成符合Conv2d要求的kernel形状 [out_channels, in_channels, kernel_H, kernel_W]
        kernel = self.l2(pre_kernel).view(1,1,5,5)

        image = x[0].view(-1, 1, 100, 100)
        # 直接调用functional接口完成卷积，无需外部Conv层
        out = self.act(F.conv2d(image, kernel, padding=2, padding_mode='replicate'))
        out = out.view(100,100)

        return out

model = Smart().to(device)
model.train()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.003)

# 优化训练循环，便于观察过程
for epoch in range(2):
    total_loss = 0.0
    for inputs, true_outputs in train_set:
        optimizer.zero_grad()
        inputs = [inp.to(device) for inp in inputs]
        true_outputs = true_outputs.to(device)

        pred_outputs = model(inputs)
        loss = criterion(pred_outputs, true_outputs)
        
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    print(f"第{epoch+1}轮训练，平均损失: {total_loss/len(train_set):.6f}")

# 验证参数是否更新
params_before = [p.data.clone() for p in model.parameters()]
# 再跑一次训练步
for inputs, true_outputs in train_set:
    optimizer.zero_grad()
    inputs = [inp.to(device) for inp in inputs]
    true_outputs = true_outputs.to(device)
    pred_outputs = model(inputs)
    loss = criterion(pred_outputs, true_outputs)
    loss.backward()
    optimizer.step()

params_after = [p.data.clone() for p in model.parameters()]
print("线性层参数是否更新:", any(not torch.equal(p1, p2) for p1, p2 in zip(params_before, params_after)))

修改说明 #

• 移除了外部的FrozenConv2d类与实例，卷积计算直接在模型内部完成
• 计算图完整覆盖从线性层生成kernel到卷积输出的全流程，梯度可正常反向传播至l1和l2
• 优化了训练循环的写法，增加epoch统计与平均损失输出，便于监控训练状态

额外建议 #

• 若后续需要支持批处理，可采用分组卷积或逐样本卷积方案（例如用torch.nn.Unfold展开图像后与每个样本的kernel做矩阵乘法）
• 可尝试调整学习率（如从0.003提升至0.01），加快参数更新速度
• 训练前可将线性层参数初始化为已知的有效权重，缩短收敛周期

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