PyTorch CNN训练中的批次大小不匹配错误：深度解析与修复

class ConvNet(nn.Module):
    def __init__(self, num_classes=4):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=4, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(in_channels=4, out_channels=8, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # 修正全连接层输入维度：16个通道 * 32x32特征图
        self.fc = nn.Linear(16 * 32 * 32, num_classes)

    def forward(self, X):
        X = F.relu(self.conv1(X))
        X = self.pool(X)
        X = F.relu(self.conv2(X))
        X = self.pool(X)
        X = F.relu(self.conv3(X))
        X = self.pool(X)
        # 扁平化操作在下一步修正
        X = X.view(X.size(0), -1) # 修正扁平化方法
        X = self.fc(X)
        return X

    def forward(self, X):
        # ... (前面的卷积和池化层保持不变)
        X = F.relu(self.conv3(X))
        X = self.pool(X)
        # 使用 X.size(0) 动态获取批次大小，-1 自动推断剩余维度
        X = X.view(X.size(0), -1)
        X = self.fc(X)
        return X

    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)

        # 修正损失函数中的标签处理，直接转换为 long 类型
        loss = criterion(outputs, labels.long())

        loss.backward()
        optimizer.step()
        # ... (其余代码保持不变)

    with torch.no_grad():
        for images, labels in val_loader:
            outputs = model(images)
            # 修正验证循环中的标签处理
            loss = criterion(outputs, labels.long())
            total_val_loss += loss.item()
            # ... (其余代码保持不变)

# ... (模型初始化、损失函数、优化器定义等)

# Placeholder for training and validation statistics
train_losses, val_losses = [], []
train_accuracies, val_accuracies = [], []

# Start training
for epoch in range(max_epoch):
    model.train() # 设置模型为训练模式
    total_train_loss = 0.0
    correct_train = 0
    total_train = 0

    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels.long()) # 修正标签处理
        loss.backward()
        optimizer.step()

        total_train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total_train += labels.size(0)
        correct_train += (predicted == labels.long()).sum().item() # 修正标签处理

    train_accuracy = correct_train / total_train if total_train > 0 else 0.0
    train_losses.append(total_train_loss / len(train_loader))
    train_accuracies.append(train_accuracy)

    # Validation
    model.eval() # 设置模型为评估模式
    total_val_loss = 0.0
    correct_val = 0 # 独立于训练的计数器
    total_val = 0   # 独立于训练的计数器

    with torch.no_grad():
        for images, labels in val_loader:
            outputs = model(images)
            loss = criterion(outputs, labels.long()) # 修正标签处理
            total_val_loss += loss.item()

            _, predicted = torch.max(outputs.data, 1)
            total_val += labels.size(0) # 正确更新验证集总数
            correct_val += (predicted == labels.long()).sum().item() # 正确更新验证集正确数，修正标签处理

    val_accuracy = correct_val / total_val if total_val > 0 else 0.0
    val_losses.append(total_val_loss / len(val_loader))
    val_accuracies.append(val_accuracy)

    print(f"Epoch {epoch+1}/{max_epoch}, "
          f"Train Loss: {train_losses[-1]:.4f}, Train Acc: {train_accuracies[-1]:.4f}, "
          f"Val Loss: {val_losses[-1]:.4f}, Val Acc: {val_accuracies[-1]:.4f}")

    # Save the best model based on validation accuracy
    if val_accuracy > best_val_accuracy:
        best_val_accuracy = val_accuracy
        best_model_state = model.state_dict()

# ... (保存模型和绘图代码)