从头实现一个自己的汽车识别分类器（中）--训练模型篇

目录：
爬取数据篇
训练模型篇
汽车识别篇

训练模型篇

数据已经爬取完毕，接下来开始训练模型，先来看一下项目结构，如下图

• all_cars：已经介绍过，存放所有类别的车型图片
• cars：其实我训练的是此文件夹数据，此文件夹包含的是从上个文件夹挑选出来的十个车型，因为直接训练十多万的数据太耗时了。所以我只选择了十个车型，大概包含一万多张图片，顺便还能解决数据不平很问题，不亏～
• test_img：是从汽车之家论坛里随便保存的几张图片，用来检验训练后模型的识别效果。
• two_cars：此文件夹存放两类车型，大概一千多张图片，用于挑选模型和初步调参之用。
• car.py和clear_data.py：爬虫和数据清理脚本，上一篇已经介绍过了。
• config .py：配置文件，用于调整超参数和更改文件路径。
• dataset.py：用于数据增强、读取等。
• main.py：核心文件，用于训练模型。
• models.py: 提供各种模型（如VGG，ResNet等）选择。
• test.py：测试模型识别效果。

首先需要介绍一下config.py，代码如下：

class Config(object):
    batch_size = 32 # 根据gpu适当增大或减小
    lr = 0.0002 # 学习率
    # lr_decay = None 学习衰减系数
    # step_size = None 学习率衰减步长
    root = './cars' # 数据集位置
    use_gpu = True # 是否使用GPU训练
    epochs = 100 # 迭代次数
    weight_decay = 0.001 # 正则化系数
    num_classes = 10 # 分类别数

这里写了一个Config类，里面设置了各种需要调解的超参数，训练模型时只需要讲此类实例化即可设置一些参数。比如学习率的设置，当然这样只是为了将超参数写一块，方便修改，也可以直接把参数直接设置到mian.py里。

opt = Config()
for epoch in opt.epochs:
      .......

接着是dataset.py文件，用于读取数据，并进行数据增强和随机打乱等操作，并为模型提供可迭代的batch_size数据。代码如下

from config import Config
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import ImageFolder
from torchvision import transforms as tfs


opt = Config()
data_tfs = tfs.Compose([tfs.Resize([160,160]),
                         tfs.RandomHorizontalFlip(),
                         tfs.RandomRotation(45),
                         tfs.ToTensor(),
                         tfs.Normalize(mean = [0.485, 0.456, 0.406],
                                        std = [0.229, 0.224, 0.225])
                        ])


data_set = ImageFolder(opt.root, transform=data_tfs)
train_size = int(0.7 * len(data_set))
val_size = len(data_set) - train_size
train_set, val_set = random_split(data_set, [train_size, val_size])

train_loader = DataLoader(train_set, batch_size=opt.batch_size, shuffle=True, num_workers=4)
val_loader = DataLoader(val_set, batch_size=opt.batch_size, shuffle=False, num_workers=4)
# for x, y in val_loader:
#     print(x,y)
#     break
print('训练集:{}, 验证集:{}'.format(len(train_loader.dataset), len(val_loader.dataset)))
print('类别：', data_set.classes)

再接着就是models.py文件，此文件很简单，就是改写一些经典的卷积神经网络，使其适应我们的分类。另外还写了一个精简版的10层VGG网络，本次用的就是次网络。

n.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)

        return x

    def _make_layers(self, cfg):
        layers = []
        in_channels = 3
        for v in cfg:
            if v == 'M':
                layers += [nn.MaxPool2d(2, 2)]
            else:
                layers += [nn.Conv2d(in_channels, v, kernel_size=3, padding=1),
                           nn.BatchNorm2d(v),
                           nn.ReLU(True)]
                in_channels = v
        return nn.Sequential(*layers)

def resnet(num_calsses):
    net = resnet(num_calsses)
    net.fc = nn.Linear(2048, num_calsses)
    # print(net)
    return net

def Alexnet(mun_classes):
    net = alexnet()
    net.classifier[6] = nn.Linear(4096, mun_classes)
    # print(net)
    return net

def squeezenet(num_classes):
    net = squeezenet1_0()
    net.classifier[1] = nn.Conv2d(512, num_classes, 1, 1)
    # print(net)
    return net

# x = torch.randn(1,3,160,160)
# net = VGG10('VGG10', 10)
# y = net(x)
# print(y.shape)

最后是训练文件mian.py，这里定义一个main函数，它对网络进行训练，然后保存一个验证集准确率最高的模型权重文件，返回训每个epoch下的训练集和测试集的loss以及准确率，并可视化。

import copy
import torch
import torch.nn as nn
from torch import optim
from config import Config
from models import VGG10, resnet, Alexnet, squeezenet
from dataset import train_loader, val_loader
import time
import os
import matplotlib.pyplot as plt

opt = Config()
device = torch.device('cuda' if opt.use_gpu else 'cpu')

net = VGG10('VGG10', opt.num_classes)
model = net.to(device)

optimizer = optim.Adam(model.parameters(), lr=opt.lr, weight_decay=opt.weight_decay)
# scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[5, 10], gamma=opt.lr_decay)
criterion = nn.CrossEntropyLoss()
train_loader = train_loader
val_loader = val_loader

def main():
    train_loss_list = []
    train_acc_list = []
    val_loss_list = []
    val_acc_list = []

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    for epoch in range(opt.epochs):
        since = time.time()
        print('epoch {}/{}'.format(epoch+1, opt.epochs))
        print('-' * 10)

        model.train()
        train_loss = 0.0
        train_corrects = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()

            _, pred = torch.max(outputs, 1)
            train_loss += loss.item() * labels.size(0)
            train_corrects += (pred == labels).sum().item()
        epoch_train_loss = train_loss / len(train_loader.dataset)
        epoch_train_acc = train_corrects / len(train_loader.dataset)

        train_loss_list.append(epoch_train_loss)
        train_acc_list.append(epoch_train_acc)

        print('train loss:{}, acc:{}'.format(epoch_train_loss, epoch_train_acc))

        model.eval()
        val_loss = 0.0
        val_corrects = 0.0
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            _, pred = torch.max(outputs, 1)

            val_loss += loss.item() * labels.size(0)
            val_corrects += (pred == labels).sum().item()
        epoch_val_loss = val_loss / len(val_loader.dataset)
        epoch_val_acc = val_corrects / len(val_loader.dataset)

        val_loss_list.append(epoch_val_loss)
        val_acc_list.append(epoch_val_acc)

        print('val loss:{}, acc:{}'.format(epoch_val_loss, epoch_val_acc))

        if epoch_val_acc > best_acc:
            best_acc = epoch_val_acc
            best_model_wts = copy.deepcopy(model.state_dict())
        time_elapsed = time.time() - since
        print('time: {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
        print()
    if not os.path.isdir('checkpoint'):
        os.makedirs('checkpoint')
        torch.save(best_model_wts, './checkpoint/weights.pkl')

    print('Best val acc:{:.4f}'.format(best_acc))

    return train_loss_list, train_acc_list, val_loss_list, val_acc_list



if __name__ == '__main__':
    train_loss_list, train_acc_list, val_loss_list, val_acc_list = main()
    plt.plot(range(1, opt.epochs+1), train_loss_list, label='train_loss')
    plt.plot(range(1, opt.epochs+1), train_acc_list, label='train_acc')
    plt.plot(range(1, opt.epochs+1), val_loss_list, label='val_loss')
    plt.plot(range(1, opt.epochs+1), val_acc_list, label='val_acc')
    plt.legend()
    plt.show()

以上就是截止到模型训练篇的所有代码。

---

分享一下训练经历吧：
由于是第一次训练自己的数据，用了好多模型，调节的好多超参数，模型要么不收敛，要么严重过拟合。再加上硬件原因，那些很深的模型很吃显存，只能把batch_size设置成16甚至更小，并且训练一个epoch都需要很长时间。于是就先提取了一个二分类的小数据集，尝试各种模型，和参数。最终resnet50和VGG，效果比较好。于是直接上resnet50，奈何训练了大半天（每个epoch大概4-5分钟，100个epoch，电脑成了暖气片～。～），严重过拟合。我感觉到一万多的数据对这种体量的网络还是太小了。那就做一个精简版的VGG吧，层数减少到10，每层通道数减半。这样可以将batch_size增加到32了，并且每个epoch缩减为1m50s，最终训练效果还是不太理想。炼丹果然是门手艺活～

附上最终绘制的loss和acc曲线。

从图可以看到，大概训练至70次左右，val_loss不再下降，说明网络还是过拟合了。从头训练一次模型实在是很耗时间和电费～
暂且就用这个模型来识别汽车吧！

图片来自汽车之家论坛