CNN模型训练循环中Batch Size不匹配问题求助

我的CNN模型（针对400x400灰度图像） #

import torch
import torch.nn as nn

class MModel(nn.Module):
    def __init__(self):
        super(MModel, self).__init__()
        
        # Define convolutional layers
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        
        # Calculate the size of the flattened feature map before the fully connected layers
        self.fc_input_size = 64 * 200 * 200
        
        # Define fully connected layers
        self.fc1 = nn.Linear(self.fc_input_size, 128)
        self.fc2 = nn.Linear(128, 18)  # Adjust the output size based on your requirements
        
    def forward(self, x):
        # Apply convolutional and pooling layers
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        
        # Flatten the feature map
        x = x.view(-1, self.fc_input_size)
        
        # Apply fully connected layers
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        
        return x

# Create an instance of the CNN model
model = MModel()

训练循环代码 #

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
num_epochs = 10

for epoch in range(num_epochs):
    running_loss = 0.0

    for i, data in enumerate(train_DL, 0):
        inputs, labels = data

        # Zero the parameter gradients
        optimizer.zero_grad()

        # Forward pass
        outputs = model(inputs)

        # Calculate the loss
        loss = criterion(outputs, labels)

        # Backward pass and optimization
        loss.backward()
        optimizer.step()

        # Print statistics
        running_loss += loss.item()

        if i % 10 == 9:  # Print every 10 mini-batches
            print(f"[{epoch + 1}, {i + 1}] Loss: {running_loss / 10:.3f}")
            running_loss = 0.0

print("Training finished")

错误信息 #

ValueError: Expected input batch_size (8) to match target batch_size (32).

设置DataLoader的batch size为32时出现上述错误，改为8仍有类似数值不匹配的错误。

解决方案 #

错误根源是手动计算的全连接层输入尺寸错误：
输入是400x400的灰度图，经过两次MaxPool2d(stride=2)，每次特征图尺寸会减半：

1. 第一次conv+pool后：400 ÷ 2 = 200（尺寸200x200）
2. 第二次conv+pool后：200 ÷ 2 = 100（尺寸100x100）

所以全连接层的输入尺寸应该是64 * 100 * 100，而不是代码中的64 * 200 * 200。错误的尺寸导致x.view(-1, self.fc_input_size)将输入batch错误地重新塑形，最终输出的batch_size和labels的batch_size不匹配。

修正方案1：手动修正尺寸计算 #

修改模型的__init__方法中的fc_input_size：

self.fc_input_size = 64 * 100 * 100

修正方案2：动态计算展平尺寸（更稳妥，避免手动计算错误） #

在forward方法中，不要手动指定展平尺寸，而是利用张量的size属性动态计算：

def forward(self, x):
    x = self.pool(nn.functional.relu(self.conv1(x)))
    x = self.pool(nn.functional.relu(self.conv2(x)))
    # 动态展平，保留batch维度，自动计算特征维度
    x = x.view(x.size(0), -1)
    x = nn.functional.relu(self.fc1(x))
    x = self.fc2(x)
    return x

同时需要修改__init__中的fc1定义，通过模拟输入计算正确尺寸：

class MModel(nn.Module):
    def __init__(self):
        super(MModel, self).__init__()
        
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        
        # 模拟输入计算全连接层输入尺寸
        with torch.no_grad():
            dummy_input = torch.randn(1, 1, 400, 400)
            x = self.pool(nn.functional.relu(self.conv1(dummy_input)))
            x = self.pool(nn.functional.relu(self.conv2(x)))
            self.fc_input_size = x.numel()
        
        self.fc1 = nn.Linear(self.fc_input_size, 128)
        self.fc2 = nn.Linear(128, 18)
        
    def forward(self, x):
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

采用任意一种修正方案后，就能解决batch_size不匹配的问题。

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