PyTorch Lightning构建CNN时线性层矩阵相乘报错求助

你的问题出在卷积池化层与线性层之间缺少张量展平操作：

卷积和MaxPool2d层输出的是4维张量（形状为(batch_size, channels, height, width)），但Linear层要求输入是2维张量（(batch_size, total_features)）。你直接在4维张量后接Linear层，导致Linear错误地将后三个维度当作特征的一部分，触发形状不匹配的RuntimeError。

解决方案 #

在最后一个MaxPool2d层之后添加nn.Flatten()层，将通道、高度、宽度维度展平为一维特征向量，再传入Linear层：

修改后的代码：

def __init__(self):    
    super().__init__()
    self.model = nn.Sequential(
        nn.Conv2d(3,6,4,padding=2),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Conv2d(6,10,4,padding=2),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Flatten(),  # 新增展平层，将4D张量转为2D
        nn.Linear(10*128*128,240),
        nn.ReLU(),
        nn.Linear(in_features=240, out_features=101),
        nn.ReLU()
    )

说明 #

nn.Flatten()默认从第1维开始展平（跳过batch维度，即索引为0的维度），会把(-1,10,128,128)转换为(-1, 10*128*128)，正好匹配Linear层的输入特征数，矩阵相乘就能正常执行。

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