PyTorch训练报错：输入batch_size(784)与目标batch_size(2)不匹配

我是PyTorch初学者，编写了针对图像数据集的简单训练与评估代码，设置batch_size=2。运行时出现错误：ValueError: Expected input batch_size (784) to match target batch_size (2)，打印的model_shape为torch.Size([2, 64, 112, 112])，代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import models, transforms

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model = models.resnet50(pretrained=True)
        self.fc1 = nn.Linear(2048,2048)
        self.fc2 = nn.Linear(2048, 3)
        self.dropout = nn.Dropout(0.3)
    def forward(self, x):
        x = torch.nn.functional.relu(self.model.conv1(x))
        print('model_shape:',x.shape)
        x = x.view(-1,2048*1*1)
        x = torch.nn.functional.relu(self.fc1(x))     
        x = F.log_softmax(self.fc2(x), dim=1)
        return x
transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
for epoch in range(100):
    running_loss = 0.0
    for i, data in enumerate(train_loader):
        inputs, labels = data
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 100 == 99:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 100))
    with torch.no_grad():
        for data in test_loader:
            inputs, labels = data
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print('Accuracy of the network on the test set: %d %%' % (100 * correct / total))

错误原因分析 #

核心问题出在forward函数的特征处理逻辑：

1. 你仅使用了ResNet50的conv1层输出，该输出形状为[2, 64, 112, 112]（batch_size=2，通道数64，特征图尺寸112x112）。
2. 错误地用x.view(-1,2048*1*1)展平特征——2048是ResNet50最后一层池化后的特征维度，并非conv1层的特征维度。计算后，展平后的张量形状变为[784, 2048]（264112*112=1605632，1605632/2048=784），此时输入的batch_size变成784，但标签的batch_size仍是2，两者维度不匹配，导致损失函数报错。

另外代码还有两个小问题：

• 训练循环未调用optimizer.zero_grad()，梯度会累加，导致训练不稳定。
• 测试循环的total和correct未初始化，每次epoch测试前未重置为0，会累计之前的结果。

解决方法 #

方法1：正确使用ResNet的完整特征提取（推荐） #

保留ResNet的完整骨干网络，仅替换最后的全连接层，这是迁移学习的标准做法：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model = models.resnet50(pretrained=True)
        # 冻结预训练骨干参数（可选，若只想训练新添加的层）
        for param in self.model.parameters():
            param.requires_grad = False
        # 获取原ResNet最后一层全连接层的输入维度
        in_features = self.model.fc.in_features
        # 替换为自定义的分类层
        self.model.fc = nn.Sequential(
            nn.Linear(in_features, 2048),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(2048, 3)
        )
    def forward(self, x):
        x = self.model(x)
        x = F.log_softmax(x, dim=1)
        return x

方法2：若坚持使用conv1层输出（不推荐，浅层特征效果差） #

需要正确计算展平后的维度，并对应修改全连接层的输入尺寸：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model = models.resnet50(pretrained=True)
        # conv1输出展平后的维度是64*112*112=784896
        self.fc1 = nn.Linear(784896,2048)
        self.fc2 = nn.Linear(2048, 3)
        self.dropout = nn.Dropout(0.3)
    def forward(self, x):
        x = torch.nn.functional.relu(self.model.conv1(x))
        # 用flatten更安全，自动计算展平后的维度
        x = torch.flatten(x, 1)
        x = torch.nn.functional.relu(self.fc1(x))
        x = self.dropout(x) # 之前定义了dropout但未使用
        x = F.log_softmax(self.fc2(x), dim=1)
        return x

修复训练和测试循环的小问题 #

for epoch in range(100):
    running_loss = 0.0
    model.train()
    for i, data in enumerate(train_loader):
        inputs, labels = data
        optimizer.zero_grad() # 新增：清零梯度
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 100 == 99:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0
    model.eval()
    total = 0 # 新增：初始化total
    correct = 0 # 新增：初始化correct
    with torch.no_grad():
        for data in test_loader:
            inputs, labels = data
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print('Accuracy of the network on the test set: %d %%' % (100 * correct / total))

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