如何修复PyTorch中输入与目标batch_size不匹配的ValueError

在使用PyTorch构建图像分类模型时，触发ValueError，提示**"Expected input batch_size (49) to match target batch_size (64)"**。调整DataLoader参数后问题仍未解决，以下是相关代码及报错堆栈：

预处理代码 #

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms


class ImageDataset(Dataset):
    def __init__(self, images, labels, transform=None, target_transform=None):
        self.images = images
        self.labels = labels
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        if self.target_transform:
            label = self.target_transform(label)

        return image, label

# 数据预处理
resize_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 构建数据集与DataLoader
train_dataset = ImageDataset(images=x, labels=y, transform=resize_transform)
train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 打印首个batch的形状
for images, labels in train_dataloader:
    print(images.shape)
    print(labels.shape)
    break

模型定义与训练代码 #

val_dataset = ImageDataset(x_val, y_val, transform=resize_transform)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)

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

class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.fc1 = nn.Linear(64 * 16 * 16, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 16)  # 适配16分类任务

        self.pool = nn.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 64 * 16 * 16)  # 扁平化层
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 定义损失函数与优化器
net = Network()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    net.train()
    running_loss = 0.0
    for i, data in enumerate(train_dataloader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)  # 报错位置
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    
    print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_dataloader):.4f}")

    # 验证集准确率计算
    net.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            outputs = net(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f"Validation Accuracy: {accuracy:.2f}%")

print("Finished Training")

完整报错堆栈 #

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[11], line 42
     40 optimizer.zero_grad()
     41 outputs = net(inputs)
---> 42 loss = criterion(outputs, labels)
     43 loss.backward()
     44 optimizer.step()

File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1736, in Module._wrapped_call_impl(self, *args, **kwargs)
   1734     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735 else:
-> 1736     return self._call_impl(*args, **kwargs)

File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1747, in Module._call_impl(self, *args, **kwargs)
   1742 # If we don't have any hooks, we want to skip the rest of the logic in
   1743 # this function, and just call forward.
   1744 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1745         or _global_backward_pre_hooks or _global_backward_hooks
   1746         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747     return forward_call(*args, **kwargs)
   1749 result = None
   1750 called_always_called_hooks = set()

File ~\anaconda3\Lib\site-packages\torch\nn\modules\loss.py:1293, in CrossEntropyLoss.forward(self, input, target)
   1292 def forward(self, input: Tensor, target: Tensor) -> Tensor:
-> 1293     return F.cross_entropy(
   1294         input,
   1295         target,
   1296         weight=self.weight,
   1297         ignore_index=self.ignore_index,
   1298         reduction=self.reduction,
   1299         label_smoothing=self.label_smoothing,
   1300     )

File ~\anaconda3\Lib\site-packages\torch\nn\functional.py:3479, in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3477 if size_average is not None or reduce is not None:
   3478     reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3479 return torch._C._nn.cross_entropy_loss(
   3480     input,
   3481     target,
   3482     weight,
   3483     _Reduction.get_enum(reduction),
   3484     ignore_index,
   3485     label_smoothing,
   3486 )

ValueError: Expected input batch_size (49) to match target batch_size (64).

问题根源与解决办法 #

1. 核心原因 #

模型forward方法中手动固定了扁平化后的特征维度：

x = x.view(-1, 64 * 16 * 16)

该写法假设经过两次卷积+池化后，特征图尺寸固定为16×16，但实际输入图像尺寸不符合这个假设时，卷积后的特征图总元素数不等于64*16*16，导致x.view强行拉伸特征，最终输出的outputs的batch_size被错误计算为49，与labels的batch_size（64）不匹配。

2. 解决方法 #

方法一：动态计算扁平化维度

不要手动写死特征尺寸，用张量属性动态获取卷积后的通道数、高度、宽度：

def forward(self, x):
    x = self.pool(F.relu(self.conv1(x)))
    x = self.pool(F.relu(self.conv2(x)))
    # 动态计算扁平化后的总维度
    x = x.view(-1, x.size(1) * x.size(2) * x.size(3))
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x

同时需要修改__init__中的fc1输入维度，可先通过一次前向传播获取实际特征数后再初始化，或根据输入图像尺寸提前计算正确的特征数（比如输入28×28图像时，两次卷积池化后特征尺寸为5×5，fc1应设为nn.Linear(64*5*5, 128)）。

方法二：统一输入图像尺寸

如果必须使用固定的64*16*16特征维度，在预处理阶段添加Resize变换，确保输入图像尺寸符合模型计算逻辑：

resize_transform = transforms.Compose([
    transforms.Resize((68, 68)),  # 确保输入尺寸经过卷积池化后得到16×16特征图
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

3. 验证步骤 #

修改后，可在训练循环中打印outputs.shape和labels.shape，确认两者的batch_size一致：

outputs = net(inputs)
print(outputs.shape)  # 应为(batch_size, 16)
print(labels.shape)   # 应为(batch_size,)

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