如何改进过拟合的PyTorch黑白图像分类TingVGG模型

我正在开展一项4类黑白图像的分类任务，训练集每类包含3000张图像，测试集总计600张图像。目前已完成模型训练，但测试效果不佳，希望了解如何改进该模型。以下是完整代码及训练结果：

数据变换定义 #

data_transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize(size=(150, 150)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5]),
])

使用ImageFolder加载图像 #

train_data = datasets.ImageFolder(root=train_dir,
                                  transform=data_transform, # Transform the data
                                  target_transform=None) # Transform the Label

test_data = datasets.ImageFolder(root=test_dir,
                                  transform=data_transform, # Transform the data
                                  target_transform=None) # Transform the Label
train_data, test_data

将数据转换为DataLoader #

BATCH_SIZE = 8

train_dataloader = DataLoader(
    dataset= train_data,
    batch_size=BATCH_SIZE, # How many images our model can see at the time
    num_workers=8,  # Number of CPU Cores
    shuffle=True
)

test_dataloader = DataLoader(
    dataset= test_data,
    batch_size=BATCH_SIZE, # How many images our model can see at the time
    num_workers=8,  # Number of CPU Cores
    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.MaxPool2d(kernel_size=2, stride=2),
        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),
        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.dropout = nn.Dropout(0.4)
        self.classifier = nn.Sequential(nn.Flatten(), nn.Linear(in_features=hidden_units*18*18 ,out_features=output_shape))
        
        
    def forward(self, x: torch.Tensor):
        x = self.conv_block1(x)
        x = self.dropout(x)
        x = self.classifier(x)
        return x

训练模型 #

# 设置随机种子
torch.manual_seed(42)
torch.cuda.manual_seed(42)

# 设置训练轮数
NUM_EPOCHS = 10

# 创建并初始化TingVGG模型
model_0 = TingVGG(input_shape=1, # 输入图像的通道数 (c, h, w) -> 3
                  hidden_units=128,
                  output_shape=len(train_data.classes)).to(device)

# 设置损失函数与优化器
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params= model_0.parameters(),
                             lr= 0.001)

# 启动计时器
start_time = time.time()

# 训练模型0
model_0_results = train(model= model_0,
                        train_dataloader= train_dataloader,
                        test_dataloader= test_dataloader,
                        optimizer= optimizer,
                        loss_fn= loss_fn,
                        epochs= NUM_EPOCHS
                        )

训练输出 #

10%|█         | 1/10 [02:30<22:37, 150.88s/it]
Epoch: 1 | train_loss: 0.3549 | train_acc: 0.8668 | test_loss: 0.3059 | test_acc: 0.8842
 20%|██        | 2/10 [04:57<19:48, 148.59s/it]
Epoch: 2 | train_loss: 0.1707 | train_acc: 0.9420 | test_loss: 0.2648 | test_acc: 0.9062
 30%|███       | 3/10 [07:24<17:14, 147.83s/it]
Epoch: 3 | train_loss: 0.1153 | train_acc: 0.9627 | test_loss: 0.2790 | test_acc: 0.8962
 40%|████      | 4/10 [09:52<14:46, 147.71s/it]
Epoch: 4 | train_loss: 0.0900 | train_acc: 0.9695 | test_loss: 0.2719 | test_acc: 0.8979
 50%|█████     | 5/10 [12:19<12:18, 147.65s/it]
Epoch: 5 | train_loss: 0.0760 | train_acc: 0.9758 | test_loss: 0.2927 | test_acc: 0.8950
 60%|██████    | 6/10 [14:47<09:50, 147.57s/it]
Epoch: 6 | train_loss: 0.0616 | train_acc: 0.9814 | test_loss: 0.3326 | test_acc: 0.8942
 70%|███████   | 7/10 [17:15<07:23, 147.76s/it]
Epoch: 7 | train_loss: 0.0488 | train_acc: 0.9838 | test_loss: 0.3086 | test_acc: 0.8946
 80%|████████  | 8/10 [19:42<04:55, 147.60s/it]
Epoch: 8 | train_loss: 0.0534 | train_acc: 0.9835 | test_loss: 0.3186 | test_acc: 0.9017
 90%|█████████ | 9/10 [22:10<02:27, 147.66s/it]
Epoch: 9 | train_loss: 0.0422 | train_acc: 0.9878 | test_loss: 0.3317 | test_acc: 0.9012
100%|██████████| 10/10 [24:38<00:00, 147.80s/it]
Epoch: 10 | train_loss: 0.0433 | train_acc: 0.9878 | test_loss: 0.3853 | test_acc: 0.9038

训练与测试的损失及准确率 #

混淆矩阵代码 #

# 导入tqdm用于进度条
from tqdm.auto import tqdm

# 1. 使用训练好的模型进行预测
y_preds = []
model_0.eval()
with torch.inference_mode():
  for X, y in tqdm(test_dataloader, desc="生成预测结果"):
    # 将数据与标签发送到目标设备
    X, y = X.to(device), y.to(device)
    # 前向传播
    y_logit = model_0(X)
    # 将预测结果从logits转为预测概率，再转为预测标签
    y_pred = torch.softmax(y_logit, dim=1).argmax(dim=1)
    # 将预测结果移至CPU用于评估
    y_preds.append(y_pred.cpu())
# 将预测结果列表拼接为张量
y_pred_tensor = torch.cat(y_preds)



from torchmetrics import ConfusionMatrix
from mlxtend.plotting import plot_confusion_matrix

# 2. 创建混淆矩阵实例并对比预测结果与真实标签
confmat = ConfusionMatrix(num_classes=len(class_names), task='multiclass')
confmat_tensor = confmat(preds=y_pred_tensor,
                         target=torch.Tensor(test_data.targets))
# 3. 绘制混淆矩阵
fig, ax = plot_confusion_matrix(
    conf_mat=confmat_tensor.numpy(), # matplotlib偏好使用NumPy格式
    class_names=class_names, # 将行列标签转为类别名称
    figsize=(10, 7)
);

混淆矩阵结果 #

从训练结果看，训练准确率接近99%但测试准确率仅约90%，且测试损失持续上升，明显存在过拟合问题，结合任务特点，可从以下几个方向优化：

1. 增强数据扩充 #

当前数据变换仅做了灰度化、Resize和归一化，缺乏数据增强，可添加针对黑白图像的变换：

data_transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize(size=(150, 150)),
    # 添加以下增强操作
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(degrees=15),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
    transforms.ColorJitter(brightness=0.2, contrast=0.2), # 针对灰度图的亮度对比度调整
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5]),
])

注意：测试集不要使用随机增强，保持原有变换即可。

2. 调整模型结构与正则化 #

（1）增加卷积块深度或调整通道数 #

当前模型只有一个卷积块，可参考VGG结构增加卷积块，同时调整通道数逐步提升：

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(input_shape, hidden_units, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.conv_block2 = nn.Sequential(
            nn.Conv2d(hidden_units, hidden_units*2, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(hidden_units*2, hidden_units*2, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.conv_block3 = nn.Sequential(
            nn.Conv2d(hidden_units*2, hidden_units*4, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(hidden_units*4, hidden_units*4, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        # 在卷积块后添加BatchNorm
        self.bn1 = nn.BatchNorm2d(hidden_units)
        self.bn2 = nn.BatchNorm2d(hidden_units*2)
        self.bn3 = nn.BatchNorm2d(hidden_units*4)
        
        self.dropout = nn.Dropout(0.5)
        # 重新计算分类器输入特征数（根据新的卷积块输出尺寸）
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(hidden_units*4*18*18, 512), # 需根据实际输出尺寸调整
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, output_shape)
        )
        
    def forward(self, x: torch.Tensor):
        x = self.conv_block1(x)
        x = self.bn1(x)
        x = self.conv_block2(x)
        x = self.bn2(x)
        x = self.conv_block3(x)
        x = self.bn3(x)
        x = self.dropout(x)
        x = self.classifier(x)
        return x

（2）添加更多正则化手段 #

• 提升Dropout率至0.5，同时在分类器中增加Dropout层
• 对卷积层和全连接层添加L2正则化（在优化器中设置weight_decay）：

optimizer = torch.optim.Adam(params=model_0.parameters(), lr=0.001, weight_decay=1e-4)

3. 调整训练策略 #

（1）使用学习率调度器 #

随着训练轮数增加，降低学习率，帮助模型收敛到更优的局部最小值：

from torch.optim.lr_scheduler import StepLR
scheduler = StepLR(optimizer, step_size=3, gamma=0.1)

# 在训练循环中添加：
for epoch in range(NUM_EPOCHS):
    # 训练步骤...
    scheduler.step()

（2）提前停止训练 #

当测试损失连续多个epoch不再下降甚至上升时，停止训练，避免过拟合：

best_test_loss = float('inf')
patience = 3
patience_counter = 0

for epoch in range(NUM_EPOCHS):
    # 训练与验证步骤，得到当前test_loss
    if test_loss < best_test_loss:
        best_test_loss = test_loss
        patience_counter = 0
        torch.save(model_0.state_dict(), 'best_model.pth')
    else:
        patience_counter +=1
        if patience_counter >= patience:
            print("提前停止训练")
            break

（3）调整批次大小 #

当前BATCH_SIZE=8过小，可尝试提升至16或32，让模型每次看到更多样本，稳定梯度更新。

4. 利用预训练模型迁移学习 #

对于图像分类任务，预训练模型在通用图像特征提取上表现更优，可针对黑白图像调整预训练模型：

from torchvision import models

# 加载预训练的VGG16，修改输入通道为1
model = models.vgg16(pretrained=True)
# 修改第一个卷积层的输入通道
model.features[0] = nn.Conv2d(1, 64, kernel_size=(3,3), stride=(1,1), padding=(1,1))
# 替换分类器头部
num_ftrs = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_ftrs, len(train_data.classes))
model = model.to(device)

冻结部分底层卷积层，仅训练顶层和分类器，减少计算量并避免过拟合。

5. 分析混淆矩阵 #

从混淆矩阵结果看，重点关注错误率高的类别，收集这些类别的样本，分析是否存在标注错误、样本特征不明显等问题，针对性补充数据或调整数据增强策略。

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