PyTorch四类图像分类模型训练异常及优化方案求助

我是PyTorch新手，正在构建一个4类图像分类模型。

编辑更新 #

将Batch_Size设置为128，并添加了新的网络层：

self.conv_block3 = nn.Sequential(nn.Conv2d(in_channels=hidden_units,out_channels=hidden_units,kernel_size=3,stride=1,padding=1),
nn.ReLU(),
nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3,stride=1,padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)

训练2个epoch后，出现异常情况，训练和测试的损失与精度如下：

Epoch: 1 | train_loss: 1.0392 | train_acc: 0.5169 | test_loss: 0.4287 | test_acc: 0.8576

Epoch: 2 | train_loss: 0.2700 | train_acc: 0.9096 | test_loss: 0.2665 | test_acc: 0.9110

想请教我的模型存在什么问题，以及该如何优化？

数据集信息 #

• 每类包含75张训练图像和25张测试图像。
• 未使用任何数据增强。

以下是输入形状及完整代码：

创建简单变换 #

simple_transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize(size=(128, 128)),
    transforms.ToTensor(),
])

train_data_simple = datasets.ImageFolder(root=train_dir,
                                          transform=simple_transform)

test_data_simple = datasets.ImageFolder(root=test_dir,
                                          transform=simple_transform)

将数据集转换为DataLoader #

BATCH_SIZE = 32
NUM_WORKERS = 2

train_dataloader_simple = DataLoader(dataset=train_data_simple,
                                     batch_size=BATCH_SIZE,
                                     num_workers=NUM_WORKERS,
                                     shuffle=True)

test_dataloader_simple = DataLoader(dataset=test_data_simple,
                                     batch_size=BATCH_SIZE,
                                     num_workers=NUM_WORKERS,
                                     shuffle=False)

创建模型类 #

class TingVGG(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()
        self.conv_block1 = nn.Sequential(nn.Conv2d(in_channels=input_shape,out_channels=hidden_units,kernel_size=3,stride=1,padding=1),
        nn.ReLU(),
        nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3,stride=1,padding=1),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        self.conv_block2 = nn.Sequential(nn.Conv2d(in_channels=hidden_units,out_channels=hidden_units,kernel_size=3,stride=1,padding=1),
        nn.ReLU(),
        nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3,stride=1,padding=1),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        self.classifier = nn.Sequential(nn.Flatten(), nn.Linear(in_features=hidden_units*32*32 ,out_features=output_shape))
        
        
    def forward(self, x: torch.Tensor):
        x = self.conv_block1(x)
        #print(x.shape)
        x = self.conv_block2(x)
        #print(x.shape)
        x = self.classifier(x)
        #print(x.shape)
        return x

创建并初始化TinyVGG模型 #

model_0 = TingVGG(input_shape=1, # Number of channels in the input image (c, h, w) -> 3
                  hidden_units=20,
                  output_shape=len(train_data.classes)).to(device)

设置损失函数与优化器 #

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params= model_0.parameters(),
                             lr= 0.001)

训练轮次、损失及精度截图：

一、当前模型的核心问题 #

1. 特征维度不匹配
   你新增了conv_block3，但原模型分类器的输入特征尺寸是hidden_units*32*32。经过三次MaxPool2d（每次尺寸减半），输入128x128的特征图最终会变成16x16，此时分类器的输入应该是hidden_units*16*16，如果没修改会直接触发维度错误，这是你所说“异常”的直接原因。

2. 小数据集+无增强，过拟合风险极高
   总训练集仅300张图（每类75张），模型2个epoch就达到90%+的训练精度，虽然当前测试精度还不错，但继续训练必然会出现训练精度持续上升、测试精度下降的过拟合现象——模型已经开始“记住”训练集的细节，而非学习通用特征。

3. 模型容量超出数据集需求
   原TinyVGG结构对300张图的数据集来说已经足够，新增conv_block3进一步增大了模型容量，会加速过拟合的发生。

二、具体优化步骤 #

1. 紧急修正维度错误

修改模型的分类器输入尺寸，并在forward函数中调用新增的conv_block3：

class TingVGG(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()
        self.conv_block1 = nn.Sequential(...)  # 保持原代码不变
        self.conv_block2 = nn.Sequential(...)  # 保持原代码不变
        # 新增conv_block3
        self.conv_block3 = nn.Sequential(nn.Conv2d(in_channels=hidden_units,out_channels=hidden_units,kernel_size=3,stride=1,padding=1),
        nn.ReLU(),
        nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3,stride=1,padding=1),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2)
        )
        # 修改分类器输入尺寸为hidden_units*16*16
        self.classifier = nn.Sequential(nn.Flatten(), nn.Linear(in_features=hidden_units*16*16 ,out_features=output_shape))
        
    def forward(self, x: torch.Tensor):
        x = self.conv_block1(x)
        x = self.conv_block2(x)
        x = self.conv_block3(x)  # 调用新增的卷积块
        x = self.classifier(x)
        return x

2. 添加数据增强，解决过拟合

针对灰度图像添加多种增强手段，大幅提升模型泛化能力：

# 训练集增强
train_transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize(size=(128, 128)),
    transforms.RandomHorizontalFlip(p=0.5),  # 随机水平翻转
    transforms.RandomRotation(degrees=15),   # 随机旋转±15度
    transforms.ColorJitter(brightness=0.2, contrast=0.2),  # 亮度对比度调整
    transforms.ToTensor(),
])

# 测试集仅做必要变换，不增强
test_transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize(size=(128, 128)),
    transforms.ToTensor(),
])

# 重新生成数据集
train_data_simple = datasets.ImageFolder(root=train_dir, transform=train_transform)
test_data_simple = datasets.ImageFolder(root=test_dir, transform=test_transform)

3. 调整训练策略与模型正则化

• 添加正则化：在卷积块或分类器中加入Dropout，或者给优化器添加权重衰减：

  # 示例：在每个卷积块的ReLU后添加Dropout2d
  self.conv_block1 = nn.Sequential(
      nn.Conv2d(...),
      nn.ReLU(),
      nn.Dropout2d(p=0.2),  # 新增Dropout
      nn.Conv2d(...),
      nn.ReLU(),
      nn.Dropout2d(p=0.2),
      nn.MaxPool2d(...)
  )
  # 优化器添加权重衰减
  optimizer = torch.optim.Adam(params=model_0.parameters(), lr=0.001, weight_decay=1e-4)
  

• 降低学习率：小数据集上0.001的学习率可能偏大，尝试调整为0.0005，或者使用学习率调度器：

  from torch.optim.lr_scheduler import ReduceLROnPlateau
  scheduler = ReduceLROnPlateau(optimizer, mode='max', patience=2, factor=0.5)
  # 每个epoch后更新学习率
  scheduler.step(test_acc)
  

• 早停机制：建议从训练集中划分10%-20%作为验证集，当验证集精度连续3个epoch不提升时，停止训练，避免过拟合。

4. 模型容量适配

如果新增conv_block3后过拟合严重，可以考虑移除该层，或者将hidden_units从20降低到16，减小模型容量以匹配小数据集。

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