PyTorch CNN训练中批次大小不匹配与维度错误：诊断与解决方案

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import os
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

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 * 32 * 32
        self.fc = nn.Linear(16 * 32 * 32, num_classes)

    def forward(self, X):
        # 卷积层、ReLU激活和最大池化
        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)

        # 展平输出，保持批次大小不变，让PyTorch自动计算其他维度
        X = X.view(X.size(0), -1)

        # 全连接层
        X = self.fc(X)

        return X

# 训练循环中
# ...
        # Forward pass
        outputs = model(images)

        # 直接使用 labels.long()，确保标签是长整型
        loss = criterion(outputs, labels.long())
# ...

# 验证循环中
# ...
        with torch.no_grad():
            for images, labels in val_loader:
                outputs = model(images)
                # 直接使用 labels.long()
                loss = criterion(outputs, labels.long())
                total_val_loss += loss.item()
# ...

# ... (其他代码保持不变，如 SceneDataset, get_dataloaders 等)

# 初始化你的网络
model = ConvNet()

# 定义你的损失函数
criterion = nn.CrossEntropyLoss()

# 初始化优化器
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, weight_decay=5e-04)

# Placeholder for best validation accuracy
best_val_accuracy = 0.0

# Placeholder for the best model state
best_model_state = None

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

# 开始训练
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).sum().item() # 修正：直接比较 predicted 和 labels

    # 计算训练准确率和损失
    train_accuracy = correct_train / total_train
    train_losses.append(total_train_loss / len(train_loader))
    train_accuracies.append(train_accuracy)

    # 验证
    model.eval() # 设置模型为评估模式
    total_val_loss = 0.0
    correct_val = 0 # 在每个epoch开始时重置
    total_val = 0   # 在每个epoch开始时重置

    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) # 修正：更新 total_val
            correct_val += (predicted == labels).sum().item() # 修正：更新 correct_val

    # 计算验证准确率和损失
    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}")

    # 根据验证准确率保存最佳模型
    if val_accuracy > best_val_accuracy:
        best_val_accuracy = val_accuracy
        best_model_state = model.state_dict()

# 保存最佳模型状态到文件
best_model_path = "best_cnn_sgd.pth"
if best_model_state:
    torch.save(best_model_state, best_model_path)
    print(f"Best model saved to {best_model_path} with validation accuracy: {best_val_accuracy:.4f}")
else:
    print("No best model saved (validation accuracy did not improve).")

# 绘制损失图
plt.figure(figsize=(10, 5))
plt.plot(train_losses, label='Training Loss')
plt.plot(val_losses, label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training and Validation Loss vs. Epoch')
plt.legend()
plt.show()

# 绘制准确率图
plt.figure(figsize=(10, 5))
plt.plot(train_accuracies, label='Training Accuracy')
plt.plot(val_accuracies, label='Validation Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Training and Validation Accuracy vs. Epoch')
plt.legend()
plt.show()