PyTorch中Country211任务神经网络损失下降缓慢问题求助

我是深度学习新手，之前基于CIFAR-10和MNIST数据集搭建过神经网络。本次尝试使用PyTorch的Country211数据集开展更大规模的分类任务，搭建了包含三层卷积层的神经网络，但训练过程中损失几乎无明显下降。不确定是训练轮数不足还是模型存在根本性问题，附上完整代码及损失输出结果，请求协助排查。

模型代码 #

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# 使用CUDA设备
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

print(device)

# 定义图像预处理（输入为PIL图像）
transform = transforms.Compose([transforms.ToTensor(), 
                                transforms.Resize((300, 300)),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

batch_size = 64

# 定义训练集和测试集
trainset = torchvision.datasets.Country211(root='./data', split='train',
                                        transform=transform, download=True)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True, num_workers=0)

testset = torchvision.datasets.Country211(root='./data', split='test',
                                          download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                        shuffle=False, num_workers=0)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 8, 5)
        self.pool = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(8, 12, 5)
        self.conv3 = nn.Conv2d(12, 16, 5)
        self.fc1 = nn.Linear(16 * 34 * 34, 4096)
        self.fc2 = nn.Linear(4096, 1024)
        self.fc3 = nn.Linear(1024, 211)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)

        return x

训练循环代码 #

net = Net()
net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
epochs = 6

for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 获取输入；data是[输入, 标签]的列表
        inputs, labels = data[0].to(device), data[1].to(device)

        # 清空参数梯度
        optimizer.zero_grad()

        # 前向传播 + 反向传播 + 优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # 打印统计信息
        running_loss += loss.item()
        if i % 2000 == 1999:    # 每2000个mini-batch打印一次
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

损失输出结果 #

[1,  2000] loss: 5.352
[1,  4000] loss: 5.351
[1,  6000] loss: 5.350
[2,  2000] loss: 5.322
[2,  4000] loss: 5.320
[3,  2000] loss: 5.276
[3,  4000] loss: 5.272
[3,  6000] loss: 5.258
[4,  2000] loss: 5.211
[4,  4000] loss: 5.197
[4,  6000] loss: 5.212
[5,  4000] loss: 5.114
[5,  6000] loss: 5.140

问题分析与解决建议 #

核心问题诊断 #

1. 模型特征提取能力不足：Country211是211类的复杂图像任务，当前模型卷积层的通道数仅为8、12、16，无法提取足够的判别特征，导致后续全连接层无法学到有效分类信息，属于欠拟合。
2. 预处理顺序错误：ToTensor先于Resize执行，虽然PyTorch支持对张量做Resize，但Resize针对PIL图像的处理更合理，顺序颠倒可能影响数据分布。
3. 学习率设置不合理：Adam优化器的初始学习率0.001对于该任务偏高，容易导致训练震荡，损失下降缓慢。
4. 缺乏训练稳定性机制：模型中没有BatchNorm或Dropout层，训练过程难以稳定收敛，也无法有效防止过拟合。

具体优化方案 #

1. 增强模型特征提取能力

增加卷积层通道数，并加入BatchNorm层提升训练稳定性：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 5)
        self.bn1 = nn.BatchNorm2d(32)
        self.pool = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 64, 5)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 128, 5)
        self.bn3 = nn.BatchNorm2d(128)
        self.fc1 = nn.Linear(128 * 34 * 34, 2048)
        self.drop1 = nn.Dropout(0.5)
        self.fc2 = nn.Linear(2048, 1024)
        self.drop2 = nn.Dropout(0.5)
        self.fc3 = nn.Linear(1024, 211)

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        x = torch.flatten(x, 1)
        x = F.relu(self.drop1(self.fc1(x)))
        x = F.relu(self.drop2(self.fc2(x)))
        x = self.fc3(x)
        return x

2. 修正预处理顺序

将Resize移至ToTensor之前：

transform = transforms.Compose([
    transforms.Resize((300, 300)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

3. 调整优化器与训练策略

• 降低初始学习率至0.0001
• 加入学习率调度器，逐步降低学习率
• 增加训练轮数至20-30轮
• 新增训练准确率监控，直观判断模型学习状态

修改后的训练循环：

net = Net()
net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.0001)
epochs = 20
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

for epoch in range(epochs):
    running_loss = 0.0
    correct = 0
    total = 0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad()

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

        # 调整打印频率，适配Country211数据集大小
        if i % 1000 == 999:
            avg_loss = running_loss / 1000
            avg_acc = 100 * correct / total
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {avg_loss:.3f}, accuracy: {avg_acc:.2f}%')
            running_loss = 0.0
            correct = 0
            total = 0
    scheduler.step()

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