干货！帮助你更好的掌握神经网络训练和模型筛选【VGG16例子、避免过拟合、可视化loss折线图、tensorboard使用例子】

前言

神经网络训练最怕的事情不是模型跑不起来，而是好不容易跑完训练之后发现模型过拟合了无法使用（血的教训），这里记录一下我自己最常用的避免过拟合的方法（并不是完全避免过拟合，只是提供了一种你可以判断是否过拟合的方式）。文章中会给出我自己编写的一个分类神经网络训练模板（GAN和目标检测的以后出），有需要的直接复制之后替换掉一些关键语句即可。下面直接上代码（我尽量都写上注释）！

引入依赖库

from tqdm import tqdm	# 进度条需要用
from datetime import datetime	# 获取时间需要用
from torchsummary import summary	# 查看神经网络结构

import torch	# pytorch框架
import torch.nn as nn	# 定义网络层需要用
import torch.optim as optim	# 定义优化函数需要用
from torchvision import datasets	# 读取数据集需要用
from torchvision import transforms	# 数据预处理需要用
from torch.utils.data import DataLoader	# 读取数据集需要用

from torch.utils.tensorboard import SummaryWriter	# 绘制loss折线图需要用

定义全局变量

这里的目录结构有一些讲究，因为做的是图像分类，我们需要将训练图像放在：/your_data_path/train/，测试图像放在/your_data_path/test/，因此这里的路径需要你自己修改一下

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"	# 是否能够使用GPU
EPOCHS = 20	# 训练的轮次
BATCH_SIZE = 16	# 数据的大小
IMAGE_SIZE = 64	# 图像的大小
LEARNINGRATE = 0.001	# 学习率
TRAIN_PATH = "datas/images/train/"	# 训练图像存放的目录（改成自己的）
TEST_PATH = "datas/images/test/"	# 测试图像存放的目录（改成自己的）
SAVE_PATH = "model/"	# 模型保存的目录
CLASSES = [0, 1]	# 分类名

定义VGG16模型

当然啦，这里你可以选择直接调用torchvision.models当中自带的其他网络，这个没关系。

class VGG16_block(nn.Module):
    def __init__(self, in_channel, out_channel, kernel, padding):
        super(VGG16_block, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=kernel, padding=padding),
            nn.BatchNorm2d(out_channel),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=out_channel, out_channels=out_channel, kernel_size=kernel, padding=padding),
            nn.BatchNorm2d(out_channel),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        x = self.layer(x)
        return x

class Vgg16Net(nn.Module):
    def __init__(self):
        super(Vgg16Net, self).__init__()

        self.layer1_4 = nn.Sequential(
            VGG16_block(in_channel=3, out_channel=64, kernel=3, padding=1),
            VGG16_block(in_channel=64, out_channel=128, kernel=3, padding=1),
            VGG16_block(in_channel=128, out_channel=256, kernel=3, padding=1),
            VGG16_block(in_channel=256, out_channel=512, kernel=3, padding=1),
        )

        self.layer5 = nn.Sequential(
            nn.Conv2d(512, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),

            nn.Conv2d(512, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),

            nn.Conv2d(512, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),

            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.fc = nn.Sequential(
            nn.Linear(512, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),

            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),

            nn.Linear(4096, 1000),
            nn.ReLU(inplace=True),
            nn.Dropout(),

            nn.Linear(1000, len(CLASSES)),
        )

    def forward(self, x):
        out = self.layer1_4(x)
        out = self.layer5(out)
        out = out.mean(dim=(2, 3))
        out = self.fc(out)
        return out

model = Vgg16Net().to(DEVICE)
summary(model, (3, IMAGE_SIZE, IMAGE_SIZE))

读取数据集

# 预处理图片
transform = transforms.Compose([
    transforms.Resize(size=(IMAGE_SIZE, IMAGE_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])

# 读取训练图片，并通过预处理
trainDataset = datasets.ImageFolder(
    root=TRAIN_PATH,
    transform=transform,
)


# 读取测试图片，并通过预处理
testDataset = datasets.ImageFolder(
    root=TRAIN_PATH,
    transform=transform,
)

# 依照batch size划分好训练集
train_dataLoader = DataLoader(dataset=trainDataset, 
                        batch_size=BATCH_SIZE, 
                        shuffle=True, 
                        num_workers=0)
                        
# 依照batch size划分好测试集
test_dataLoader = DataLoader(dataset=testDataset, 
                        batch_size=BATCH_SIZE, 
                        shuffle=True, 
                        num_workers=0)

print('Train dataSize: {}\nTrain label: {}'.format(len(train_dataLoader.dataset), trainDataset.classes))
print('Test dataSize: {}\nTest label: {}'.format(len(test_dataLoader.dataset), testDataset.classes))

定义损失函数和优化函数（这里可以选择适合自己神经网络的）

loss_func = nn.CrossEntropyLoss()	# 损失函数
optimizer = optim.SGD(model.parameters(), lr=LEARNINGRATE)	 # 优化函数

# 创建一份 tensorboard 文件，存放在logs文件夹下（这个路径可以自己调整）
writer_tensorboard = SummaryWriter(f"logs/" + str(datetime.now().strftime("%Y%m%d_%H%M%S")) + "/") 

开始训练（这里的代码有点多，如果看不懂的话可以复制到GPT问问）

SAVE_MODEL = 98.00	# 保存模型的阈值

for epoch in range(EPOCHS):
    with tqdm(total=len(train_dataLoader), desc="Epoch {}".format(epoch)) as trainpbar:
        model.train()
        train_loss = 0
        Train_Accuracy = 0
        for step, train_data in enumerate(train_dataLoader):
            images, labels = train_data
            images = images.to(DEVICE)
            output = model(images)

            _, pre_tensor = torch.max(input=output, dim=1)
            if torch.eq(pre_tensor, labels.to(DEVICE)).all():
                    Train_Accuracy += 1
            
            optimizer.zero_grad()
            loss = loss_func(output, labels.to(DEVICE))
            loss.backward(retain_graph=True)
            optimizer.step()

            train_loss += loss.item()

            trainpbar.update(1)
            trainpbar.set_postfix({"Train_Loss": loss.item() / (step + 1),
                                   "Train_Accuracy": Train_Accuracy/len(train_dataLoader)*100})
        writer_tensorboard.add_scalar(tag="Train Loss", scalar_value=train_loss/len(train_dataLoader), global_step=epoch)
        writer_tensorboard.add_scalar(tag="Train Accuracy", scalar_value=Train_Accuracy/len(train_dataLoader)*100, global_step=epoch)

    if (epoch+1) % 11 == 0:
        model.eval()
        Test_Accuracy = 0
        with tqdm(total=len(test_dataLoader), desc="Test {}".format(epoch)) as testpbar:
            for step, test_data in enumerate(test_dataLoader):
                images, labels = test_data
                _, pre_tensor = torch.max(input=model(images.to(DEVICE)), dim=1)

                if torch.eq(pre_tensor, labels.to(DEVICE)).all():
                    Test_Accuracy += 1
                testpbar.update(1)
                testpbar.set_postfix({"Test_Accuracy": Test_Accuracy/len(test_dataLoader)*100})

        Test_Accuracy = Test_Accuracy/len(test_dataLoader)*100
        writer_tensorboard.add_scalar(tag="Test Accuracy", scalar_value=Test_Accuracy, global_step=epoch)
        
        if Test_Accuracy > SAVE_MODEL:
            torch.save(model.state_dict(), SAVE_PATH + str(epoch) + "_" + str(Test_Accuracy) + ".pt")
            SAVE_MODEL = Test_Accuracy

结语

在训练的时候可以在logs文件夹（也就是定义损失函数和优化函数这一步当中定义的）下打开命令行，输入tensorboard --logdir=logs命令，打开一个网页，在这个网页的右上角有一个小齿轮，点进去勾选一下第一个选项，这样每隔30s会自动刷新一次数据，就可以看到数据变化啦！

那么如何避免过拟合呢？

1. 数据集扩充（Data Augmentation）：通过对训练数据进行随机变换，如旋转、缩放、平移等，可以生成更多的训练样本。这样有助于增加训练数据的多样性，提高模型的泛化能力。
2. 正则化（Regularization）：正则化是通过在损失函数中引入额外的惩罚项来避免模型过拟合。常见的正则化方法包括L1正则化和L2正则化。正则化通过减小模型的复杂度来控制过拟合。
3. 早停（Early Stopping）：早停是在训练过程中监控验证集上的性能，并在性能不再提升时停止训练。这样可以防止模型在训练数据上过拟合。
4. Dropout：Dropout是一种正则化技术，在训练过程中随机地丢弃一些神经元的输出。这阻止了神经元之间过度依赖的情况，增强了模型的泛化能力。
5. 权重衰减（Weight Decay）：权重衰减通过在损失函数中添加权重的L2范数惩罚项，减小权重的值，从而降低模型的复杂度。
6. 模型复杂度控制：通过调整模型的复杂度，如减少网络层数、单元数等，可以减小模型在训练数据上的拟合程度。
7. 批标准化（Batch Normalization）：批标准化在每一层的输入数据上进行归一化处理，有助于防止模型在训练过程中的内部协变量偏移，从而提高模型的泛化能力。

我这里提供的方法就属于早停法，我们可以根据tensorboard展现出来的数据折线图，推断这个模型的准确率和loss值是不是已经不再增长，或者早在十几个epoch之前就已经到达了准确最大loss最小，这时候就很有可能出现过拟合的情况，应当停止训练。如果这时候的模型使用的时候表现得不是很好，那么就要考虑优化模型了。

这里附上我的结果图：