训练模型篇
数据已经爬取完毕,接下来开始训练模型,先来看一下项目结构,如下图
-
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不再下降,说明网络还是过拟合了。从头训练一次模型实在是很耗时间和电费~
暂且就用这个模型来识别汽车吧!