PyTorch自定义类VGG模型仅batch_size=1可用，其他值报尺寸不匹配错误

算力阿峰

2026-6-28

问题：Batch Size≠1时模型训练报错

使用PyTorch实现类VGG19的神经网络模型，当设置batch_size为1以外的值时，抛出错误：

ValueError: Expected input batch_size (1) to match target batch_size (16).

仅当batch_size=1时模型能正常训练，希望修改batch_size且不触发错误，以下是模型代码及完整错误回溯：

模型代码

import os
import glob

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from torchvision import datasets, transforms

class ICModel(nn.Module):

    def __init__(self):
        super().__init__()
        # CNN
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=2)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=2)
        self.conv4 = nn.Conv2d(256, 512, kernel_size=3, stride=2)

        # FULLY CONNECTED LAYERS
        self.fc1 = nn.Linear(61952, 256)
        self.fc2 = nn.Linear(256, 64)
        self.out = nn.Linear(64, 2)
    def forward(self, x):
        # CONV - 1
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, kernel_size=3, stride=1)
        # CONV - 2
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, kernel_size=3, stride=1)
        # CONV - 3
        x = F.relu(self.conv3(x))
        x = F.max_pool2d(x, kernel_size=3, stride=1)
        # CONV - 4
        x = F.relu(self.conv4(x))

        flattened_size = x.shape[0] * x.shape[1] * x.shape[2] * x.shape[3]

        x = x.view(-1, flattened_size)
        # FULLY CONNECTED LAYERS
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.out(x)

        return F.log_softmax(x, dim=1)

    def train(self, dataset_dir='', epochs=5, batch_size=16, seed=35, learning_rate=0.001, model_weights_path=''):
        if dataset_dir == '':
            raise Exception("Please enter a valid dataset directory path!")

        train_correct = []
        train_losses = []

        torch.manual_seed(seed)

        # CRITERION AND OPTIMIZER SETUP
        criterion = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)

        optim_width, optim_height = 224, 224

        data_transforms = transforms.Compose([
            transforms.Resize((optim_width, optim_height)),  # Resize images to average dimensions
            transforms.ToTensor(),  # Convert images to PyTorch tensors
            transforms.Normalize(mean=[0.456, 0.456, 0.456], std=[0.456, 0.456, 0.456])  # Normalize images
        ])

        dataset = datasets.ImageFolder(root=dataset_dir, transform=data_transforms)
        train_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True)

        for epoch in range(epochs):
            trn_corr = 0

            for b, (X_train, y_train) in enumerate(train_loader):
                b += 1
                y_pred = self(X_train)
                loss = criterion(y_pred, y_train)

                predicted = torch.max(y_pred, dim=1)[1]
                batch_corr = (predicted == y_train).sum()

                trn_corr += batch_corr.item()

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                if b % 4 == 0:
                    print(f'Epoch: {epoch}  Batch: {b}  Loss: {loss.item()}')

        train_losses.append(loss)
        train_correct.append(trn_corr)

        if (model_weights_path != '') & os.path.exists(model_weights_path) & os.path.isdir(model_weights_path):
            torch.save({
                'model_state_dict': self.state_dict(),
                'optimizer_state_dict': optimizer.state_dict()
            }, model_weights_path)

    def test(self, dataset_dir='', batch_size=16):
        if dataset_dir == '':
            raise Exception("Please enter a valid dataset directory path!")

        optim_width, optim_height = 224, 224
        test_losses = []
        tst_crr = 0

        criterion = nn.CrossEntropyLoss()

        data_transforms = transforms.Compose([
            transforms.Resize((optim_width, optim_height)),  # Resize images to average dimensions
            transforms.Grayscale(),
            transforms.ToTensor(),  # Convert images to PyTorch tensors
            transforms.Normalize(mean=[0.456], std=[0.456])  # Normalize images
        ])

        dataset = datasets.ImageFolder(root=dataset_dir, transform=data_transforms)
        test_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True)

        with torch.no_grad():
            for b, (X_test, y_test) in enumerate(test_loader):
                y_val = self(X_test)
                predicted = torch.max(y_val.data, dim=1)[1]
                tst_crr += (predicted == y_test).sum()

            loss = criterion(y_val, y_test)
            test_losses.append(loss.item())

            test_results = {
                'true_positive': tst_crr,
                'false_positive': len(dataset.imgs) - tst_crr
            }

        return test_results, test_losses

错误回溯

Traceback (most recent call last):
  File "/Users/eaidy/Repos/ML/inclination-classification-pytorch/src/main.py", line 12, in <module>
    ic_model.train(dataset_dir=train_dataset_absolute_path, epochs=1, batch_size=16, learning_rate=1e-5)
  File "/Users/eaidy/Repos/ML/inclination-classification-pytorch/src/models/cnn_model.py", line 94, in train
    loss = criterion(y_pred, y_train)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/eaidy/Repos/ML/inclination-classification-pytorch/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/eaidy/Repos/ML/inclination-classification-pytorch/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/eaidy/Repos/ML/inclination-classification-pytorch/.venv/lib/python3.11/site-packages/torch/nn/modules/loss.py", line 1179, in forward
    return F.cross_entropy(input, target, weight=self.weight,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/eaidy/Repos/ML/inclination-classification-pytorch/.venv/lib/python3.11/site-packages/torch/nn/functional.py", line 3059, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: Expected input batch_size (1) to match target batch_size (16).

解决方案

错误根源在forward函数的特征展平步骤：

flattened_size = x.shape[0] * x.shape[1] * x.shape[2] * x.shape[3]
x = x.view(-1, flattened_size)

这里的x.shape[0]是batch大小，计算flattened_size时把batch维度也乘进去，导致每个batch被压缩成1个样本（总元素数为batch_size通道数高*宽，view(-1, flattened_size)就会将整个batch变成1行，对应batch_size=1），和输入的target batch_size（比如16）不匹配，因此报错。

修改方法有两种，任选其一即可：

方法1：修正flattened_size的计算

排除batch维度，只计算单样本的特征维度：

def forward(self, x):
    # CONV - 1
    x = F.relu(self.conv1(x))
    x = F.max_pool2d(x, kernel_size=3, stride=1)
    # CONV - 2
    x = F.relu(self.conv2(x))
    x = F.max_pool2d(x, kernel_size=3, stride=1)
    # CONV - 3
    x = F.relu(self.conv3(x))
    x = F.max_pool2d(x, kernel_size=3, stride=1)
    # CONV - 4
    x = F.relu(self.conv4(x))

    # 只计算单样本的特征维度，排除batch维度x.shape[0]
    flattened_size = x.shape[1] * x.shape[2] * x.shape[3]
    x = x.view(-1, flattened_size)
    # FULLY CONNECTED LAYERS
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.out(x)

    return F.log_softmax(x, dim=1)

方法2：使用更简洁的展平方式

直接保留batch维度，展平后续所有维度，无需手动计算尺寸：

def forward(self, x):
    # CONV - 1
    x = F.relu(self.conv1(x))
    x = F.max_pool2d(x, kernel_size=3, stride=1)
    # CONV - 2
    x = F.relu(self.conv2(x))
    x = F.max_pool2d(x, kernel_size=3, stride=1)
    # CONV - 3
    x = F.relu(self.conv3(x))
    x = F.max_pool2d(x, kernel_size=3, stride=1)
    # CONV - 4
    x = F.relu(self.conv4(x))

    # 保留batch维度（dim=0），展平从dim=1开始的所有维度
    x = x.flatten(1)
    # 或者等价写法：x = x.view(x.size(0), -1)
    # FULLY CONNECTED LAYERS
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.out(x)

    return F.log_softmax(x, dim=1)

修改后，模型输出的batch_size会和输入一致，即可正常使用任意batch_size训练。

另外注意：测试函数test中存在两个潜在问题：

测试时转换用了transforms.Grayscale()，将3通道转为单通道，但模型输入是3通道（conv1的in_channels=3），会导致输入维度不匹配，需删除该转换；
损失计算放在循环外，只会计算最后一个batch的损失，需移到循环内或修改逻辑。

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

PyTorch自定义类VGG模型仅batch_size=1可用，其他值报尺寸不匹配错误

模型代码 #

错误回溯 #

方法1：修正flattened_size的计算 #

方法2：使用更简洁的展平方式 #

模型代码

错误回溯

方法1：修正flattened_size的计算

方法2：使用更简洁的展平方式