PyTorch中CNN植物病害预测模型矩阵相乘维度不匹配RuntimeError问题求助

Hey there! Let's break down what's causing this error and how to fix it quickly.

The Root Cause #

That RuntimeError spells out the issue clearly: the flattened output from your convolutional layers (shape 32x119072) doesn't match the input expected by your first fully connected layer (800x300).

When nn.Flatten() runs after your CNN block, it converts the 3D feature map (channels × height × width) into a 1D vector. Your hardcoded nn.Linear(800, 300) assumes this vector is 800 elements long, but for your input images, it's actually 119072 elements long—hence the mismatch.

How to Fix It #

You have two solid approaches here: one that's manual, and one that's dynamic (so you don't have to recalculate if you tweak your model or input size later).

Option 1: Dynamically Calculate the Flattened Dimension (Recommended)

This method uses a dummy input to automatically figure out the correct size of the flattened vector after the convolutional layers. Here's how to adjust your code:

import torch
import torch.nn as nn

class CNN_Model(nn.Module):
    def __init__(self, input_shape):
        super(CNN_Model, self).__init__()
        self.cnn_model = nn.Sequential(
            nn.Conv2d(3, 16, 3),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(16, 32, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
        )
        
        # Use a dummy input to compute the correct flattened size
        with torch.no_grad():
            # Dummy input matches your actual input shape: (channels, height, width)
            dummy_input = torch.randn(1, *input_shape)
            cnn_output = self.cnn_model(dummy_input)
            # Calculate elements per sample (divide by batch size=1)
            flattened_dim = cnn_output.numel() // dummy_input.size(0)
        
        self.fc_model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(flattened_dim, 300),
            nn.ReLU(),
            nn.Linear(300, 38),
            nn.Softmax(dim=1)
        )
    
    def forward(self, x):
        x = self.cnn_model(x)
        x = self.fc_model(x)
        return x

# Initialize with your actual input image shape (e.g., 3 channels, 256x256)
# Based on your error, your input is likely 256x256 (32*61*61=119072)
model = CNN_Model(input_shape=(3, 256, 256))

# Test with a sample batch
imgs = torch.randn(32, 3, 256, 256)  # Batch size 32, matching your input shape
out = model(imgs)
print(out.shape)  # Should output torch.Size([32, 38])

Option 2: Manually Calculate the Flattened Dimension

If you prefer doing the math yourself, here's how to work backwards from your error:

1. Your error says the flattened shape is 32x119072—32 is the channel count from the last conv layer, so 119072 / 32 = 3721, which is 61×61 (the height/width of the final feature map).
2. Reverse-engineer the input size:
   • After the second MaxPool (2x2), you have 61x61 → before pooling, it was 123x123 (since 123//2=61).
   • The second Conv2d uses a 5x5 kernel: input size before this conv was 123 +5 -1 =127x127.
   • After the first MaxPool (2x2), you have 127x127 → before pooling, it was 254x254.
   • The first Conv2d uses a 3x3 kernel: original input size is 254 +3 -1=256x256.
3. Update your fully connected layer to use the correct flattened size:
   Change nn.Linear(800, 300) to nn.Linear(32*61*61, 300) (or directly nn.Linear(119072, 300)).

Quick Notes #

• Make sure all your input images are resized to the same dimensions—shape mismatches will trigger this error again.
• The dynamic method (Option 1) saves you time if you later adjust your CNN layers (e.g., change kernel sizes, add more conv blocks).

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