PyTorch中LazyLinear全连接层无参数问题排查求助

问题原因 #

nn.LazyLinear的参数是延迟初始化的——只有当它第一次处理输入张量时，才会根据输入维度自动创建权重和偏置参数。你看到参数为0，是因为打印参数时这个层还没被调用过，参数尚未生成。

同时还有两个影响训练的关键问题：

• CrossEntropyLoss内部已经集成了Softmax计算，你额外添加的nn.Softmax会导致损失计算逻辑错误，直接让网络无法正常学习。
• 优化器的学习率0.05过高，RMSprop在这个学习率下极易出现训练不稳定、损失发散的情况。

解决方案 #

方案1：手动计算维度，替换LazyLinear（更稳妥）

先根据输入图像尺寸计算卷积层输出后的扁平化维度（以MNIST的28×28单通道图像为例）：

1. 第一次Conv+MaxPool：28×28 → 14×14（通道数16）
2. 第二次Conv+MaxPool：14×14 →7×7（通道数32）
3. 第三次Conv后：7×7（通道数64）
   扁平化后的维度为 64*7*7=3136，直接用固定维度的nn.Linear替代LazyLinear。

修改后的核心代码：

# __init__中替换全连接层定义
self.fc1 = nn.Linear(64*7*7, num_classes)  # 替换原LazyLinear
# 移除softmax层，CrossEntropyLoss自带该逻辑
# self.softmax = nn.Softmax(dim=1)

# forward方法中移除softmax调用
def forward(self, x):
    # ... 保留原卷积、池化、归一化步骤 ...
    out = out.reshape(out.size(0), -1)
    out = self.fc1(out)
    # 移除 out = self.softmax(out)
    return out

方案2：强制触发LazyLinear初始化

如果坚持使用LazyLinear，可以在初始化时传入一个dummy张量触发前向传播，让参数提前生成：

def __init__(self, num_classes, num_images):
    super(ThreeConvLayer, self).__init__()
    # ... 保留所有层的定义 ...
    # 用dummy张量触发LazyLinear初始化
    with torch.no_grad():
        dummy_input = torch.randn(1, 1, 28, 28)  # 匹配你的输入尺寸
        self(dummy_input)

同样需要移除Softmax层，并调低学习率。

额外修正点 #

1. 训练日志命名错误：training_step中你log的是val_loss，但这是训练步骤，应改为train_loss避免混淆：

def training_step(self, batch, batch_no):
    x, y = batch
    y = y.long()
    logits = self(x)
    loss = self.loss(logits, y)
    self.log('train_loss', loss, prog_bar=True)
    return loss

2. 优化器学习率调整：将RMSprop的学习率调低到0.001或0.0005：

def configure_optimizers(self):
    return torch.optim.RMSprop(self.parameters(), lr=0.001)

3. 冗余参数清理：num_images参数在代码中未被使用，可直接移除。

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