Keras转PyTorch模型报错：矩阵形状无法相乘问题求助

我用Keras构建的CNN模型训练正常，但转成PyTorch版本运行时触发错误：
RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x4851 and 155232x128)

我完全沿用了Keras的参数设置，但还是出现这个问题，以下是相关代码和模型信息：

原Keras代码 #

input_shape = 200, 100, 1

def make_model():
    return keras.models.Sequential([
        keras.layers.Conv2D(32, kernel_size=(2, 2), activation='relu', input_shape=input_shape),
        keras.layers.MaxPooling2D(pool_size=(2, 2)),
        keras.layers.Dropout(0.2),
        keras.layers.Flatten(),
        keras.layers.Dense(128, activation='relu'),
        keras.layers.Dense(64, activation='relu'),
        keras.layers.Dense(3, activation='sigmoid'),
    ])
    
model = make_model()

Keras模型摘要 #

Model: "sequential_164"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 conv2d_16 (Conv2D)          (None, 199, 99, 32)       160       
                                                                 
 max_pooling2d_16 (MaxPooli  (None, 99, 49, 32)        0         
 ng2D)                                                           
                                                                 
 dropout_66 (Dropout)        (None, 99, 49, 32)        0         
                                                                 
 flatten_23 (Flatten)        (None, 155232)            0         
                                                                 
 dense_322 (Dense)           (None, 128)               19869824  
                                                                 
 dense_323 (Dense)           (None, 51)                6579      
                                                                 
=================================================================
Total params: 19876563 (75.82 MB)
Trainable params: 19876563 (75.82 MB)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________

我的PyTorch尝试代码 #

class MyCNN(nn.Module):
    def __init__(self):
        super(CustomCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=(2, 2))
        self.pool = nn.MaxPool2d(2)
        self.dropout = nn.Dropout(0.2)
        self.fc1 = nn.Linear(155232, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 3)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = self.dropout(x)
        x = torch.flatten(x, start_dim=1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.softmax(self.fc3(x))
        return x

错误原因 #

问题出在三个核心点：

1. 输入维度顺序差异：Keras默认输入维度是(batch_size, height, width, channels)，但PyTorch默认是(batch_size, channels, height, width)。如果直接传入Keras格式的输入，会被PyTorch误判通道数，导致卷积池化后的特征形状完全错误。
2. 类名初始化错误：super(CustomCNN, self).__init__()中的类名和定义的MyCNN不匹配，会引发初始化逻辑异常。
3. 激活函数不一致：原Keras模型输出层用sigmoid，但PyTorch代码用了softmax，不仅改变输出逻辑，也可能间接影响张量处理流程。

修正后的PyTorch代码 #

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

class MyCNN(nn.Module):
    def __init__(self):
        super(MyCNN, self).__init__()
        # 对齐Keras的卷积参数：输入通道1，输出32，卷积核(2,2)
        self.conv1 = nn.Conv2d(1, 32, kernel_size=(2, 2))
        self.pool = nn.MaxPool2d((2, 2))
        self.dropout = nn.Dropout(0.2)
        # 计算池化后的总特征数：32 * 99 * 49 = 155232，和Keras一致
        self.fc1 = nn.Linear(32 * 99 * 49, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 3)

    def forward(self, x):
        # 若输入是Keras格式(batch, H, W, C)，需先转置为PyTorch格式：x = x.permute(0,3,1,2)
        x = F.relu(self.conv1(x))  # 输出形状：(batch, 32, 199, 99)
        x = self.pool(x)           # 输出形状：(batch, 32, 99, 49)
        x = self.dropout(x)
        x = torch.flatten(x, start_dim=1)  # 展平为(batch, 155232)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = torch.sigmoid(self.fc3(x))  # 和Keras保持一致用sigmoid
        return x

# 测试输入匹配性
if __name__ == "__main__":
    model = MyCNN()
    # PyTorch标准输入形状：(batch_size, channels, height, width)
    test_input = torch.randn(32, 1, 200, 100)
    output = model(test_input)
    print(output.shape)  # 应输出torch.Size([32, 3])

关键注意事项 #

• 输入维度转换：如果你的数据集是Keras格式，必须用x.permute(0,3,1,2)转成PyTorch要求的维度顺序。
• 类名一致性：super函数中的类名必须和定义的类名完全一致。
• 激活函数对齐：输出层激活函数要和原模型保持一致，避免任务逻辑偏差。

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