Pytorch物体识别分类实战案例(下)
前一篇文章介绍了案例中关于数据集的部分,接下来这部分是最重要的部分,包括网络的搭建、loss函数、训练和测试,还是话不多说,直接看代码。
5.vgg16_net.py:搭建vgg16网络
该py文件主要用于搭建网络结构,如果需要自己搭建不同的网络模型,直接仿造本文件去编写,先更改py文件的名字,定义网络的类名,在__init__f方法中对网络层进行定义,在forward方法中构建前向传播过程,大家仔细品一品这个流程,以后搭建自己的网络模型就是得心应手了!
本项目搭建的是最经典的vgg16网络,先看vgg16的网络结构图:
1.224x224x3的彩色图表示3通道的长和宽都为224的图像数据,也是网络的输入层
2.白色部分为卷积层,红色部分为池化层(使用最大池化),蓝色部分为全连接层,其中卷积层和全连接层的激活函数都使用relu
3.总的来说,VGG16网络为13层卷积层+3层全连接层而组成,从图中可以看出大概分为了5个block,这与代码中的block一一对应,对着16的网络层就可以搭建了,但是得仔细,参数容易写错。
import torch.nn as nn #利用nn模块进行神经网络的搭建
import config as cfg
class Vgg16_Net(nn.Module): #定义网络结构的类名
def __init__(self): #初始化参数
super(Vgg16_Net, self).__init__() #继承定义的类
self.num_classes = cfg.NUM_CLASSES #定义类别数,在config.py中已定义
#vgg_16的网络结构
#block1
self.block1 = nn.Sequential( #nn.Sequential用于连接该函数里面所有的网络层
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1), #定义卷积层
nn.ReLU(), #激活函数层
nn.Conv2d(64, 64, 3, 1, 1), #卷积层
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2) #最大池化层
)
#block2
self.block2 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(128, 128, 3, 1, 1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
#block3
self.block3 = nn.Sequential(
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(256, 256, 3, 1, 1),
nn.ReLU(),
nn.Conv2d(256, 256, 3, 1, 1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
#block4
self.block4 = nn.Sequential(
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 512, 3, 1, 1),
nn.ReLU(),
nn.Conv2d(512, 512, 3, 1, 1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
#block5
self.block5 = nn.Sequential(
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(512, 512, 3, 1, 1),
nn.ReLU(),
nn.Conv2d(512, 512, 3, 1, 1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
#平均池化
self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
#第一个全连接层
self.fc1 = nn.Sequential(
nn.Linear(512*7*7, 4096),
nn.ReLU(),
nn.Dropout()
)
#第二个全连接层
self.fc2 = nn.Sequential(
nn.Linear(4096, 4096),
nn.ReLU(),
nn.Dropout()
)
#第三个全连接层,也是最终的预测层
self.classifier = nn.Linear(4096, self.num_classes)
def forward(self, x): #定义前向传播的过程
x = self.block1(x)
x = self.block2(x)
x = self.block3(x)
x = self.block4(x)
x = self.block5(x)
x = self.avgpool(x)
x = x.view(x.size(0), -1) #进入全连接层时需要将输出拉成一行
x = self.fc1(x)
x = self.fc2(x)
x = self.classifier(x) #到这就得到最终的预测结果
return x
def _initialize_weights(self): #定义权重初始化的函数
for m in self.modules(): #遍历所有的层
if isinstance(m, nn.Conv2d): #对卷积层的初始化方案
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d): #对BN层的初始化方案
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear): #对全连接层的初始化方案
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
6.loss_function.py:损失函数
本案例只是最简单的二分类任务,因此使用的损失函数也就是一个交叉熵的损失,比较简单,但是为了理解torch的使用和代码结构,将其单独写进了loss_function.py中,用于学习整个框架。
每个网络都会有自己单独的损失函数,大家要是使用了不同的paper中的网络模型,只需要找到对应的损失函数的部分代码,将其copy到该文件下重写loss即可,或者也可以自己重新编写。
import torch.nn as nn
class Loss(nn.Module): #定义loss的类,也是继承自nn.Module
def __init__(self): #初始化
super(Loss,self).__init__() #继承Loss类
def forward(self, predictions, targets): #定义损失的前向传播方法
cls_pred = predictions #拿到预测label
cls_true = targets #拿到真实的label
criterion = nn.CrossEntropyLoss() #初始化交叉熵损失的类
loss = criterion(cls_pred, cls_true.long()) #交叉熵损失函数,一定得把预测结果放第一位,真实值放第二位
return loss
7.train.py 训练代码
训练部分都是大同小异,先读入数据,定义好model,优化器optimizer,损失函数,然后遍历epoch,计算损失,反向传播,更新参数,保存模型,也算是一个标准的模板,大家好好感受一下整个过程就明白了。
本部分代码其中包含了torch的断点续训,gpu的调用,学习率的调整,基本算是一个很完整的模板了。
import config as cfg
import torch
from vgg16_net import Vgg16_Net
from data_lodar import Data_loader
from torchvision import transforms
from data_augmentation import Rescale,RandomCrop,ToTensor
from torch.utils.data import DataLoader
from loss_function import Loss
from torch.optim.lr_scheduler import ReduceLROnPlateau
'''
自定义学习率调整策略,也可以不自定义,torch里有许多策略,本案例就使用的torch里已有的
'''
# def adjust_lr(optimizer, lr, gamma, step):
# lr = lr * (gamma ** (step))
# for param_group in optimizer.param_groups:
# param_group['lr'] = lr
# return lr
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #若gpu可用就用gpu训练,反之cpu训练
lr = cfg.LEARNING_RATE #学习率
epochs = cfg.EPOCHS #数据集训练次数
batch_size = cfg.BATCH_SIZE #batch的大小
resume = './model/Eopch1-Step1-cls_loss0.6872.pth' #是否断点续训或者加载预训练权重
# resume = None
train_file = 'train.txt' #训练数据
with open(train_file) as f: #读取训练数据
lines = f.readlines()
#生成自定义的Data_loader生成数据集
train_dataset = Data_loader(lines,
transform=transforms.Compose([ #定义数据增强方法,Compose就是一个连接操作
Rescale(256), #随机缩放
RandomCrop(), #随机裁剪
ToTensor() #转换为tensor
]))
#通过torch里的DataLoader类来加载一个batch的数据集
dataloader = DataLoader(train_dataset, batch_size=batch_size,
shuffle=True) #shuffle操作为打乱顺序,能增强模型的鲁棒性
model = Vgg16_Net() #定义model结构
net = model.to(device) #将model防御gpu或者cpu上
loss_fun = Loss() #定义loss函数
if resume: #是否断点续训或加载预训练权重
checkpoint = torch.load(resume) #使用torch.load加载模型
net.load_state_dict(checkpoint['model']) #加载权重
optimizer = checkpoint['optimizer'] #加载优化器
start_epoch = checkpoint['epoch'] #加载断点的epoch
#print(start_epoch)
print('加载epoch{}成功!'.format(start_epoch))
else: #不断点续训
start_epoch = 0 #从0开始训练
print('无保存模型,将从头开始训练!')
optimizer = torch.optim.SGD(net.parameters(), lr=lr) #定义优化器
epoch_size = len(lines) // batch_size #epoch_size为一个epoch中需要训练的step数
#比如,数据总量为100,bath=2
#那么整个数据集训练一次的步数为50
#学习率的优化策略,具体方法下面再细说
scheduler = ReduceLROnPlateau(optimizer, mode="min", factor=0.1,
patience=15, verbose=True, threshold=0.00001,
threshold_mode='rel',cooldown=0, min_lr=0,
eps=1e-08)
#print('start_epoch', start_epoch)
for epoch in range(start_epoch, epochs): #循环epochs次
class_loss = 0 #总的loss
for step, batch_sample in enumerate(dataloader): #循环一个epoch中的所有step的数据
image = batch_sample['image'] #取出一个batch中的图片
label = batch_sample['label'] #取出一个batch中对应的label
'''Variable方法在老版本的torch中用于构建torch需要的变量,
只有Variable变量才可以反向传播,在新版的torch中已经被弃用,
不再需要单独构建Variable,tensor直接就可以进行反向传播,git
上的许多开源代码依然使用了Variable,大家可以看着更改'''
# image = Variable(image.type(torch.FloatTensor)).to(device)
# label = Variable(label.type(torch.FloatTensor)).to(device)
image = image.float() #将image转化为float型
image.to(device) #将image传入gpu或者cpu
label.to(device) #将label传入gpu或者cpu
pred = net(image) #前向传播得到预测值
cls_loss = loss_fun(pred, label) #传入损失函数计算损失
optimizer.zero_grad() #优化器重置梯度为0,这步必须有
cls_loss.backward() #对loss进行反向传播更新
optimizer.step() #更新优化器参数
class_loss = class_loss + cls_loss.item() #计算总损失,.item()用于查看tensor的值
scheduler.step(class_loss) #以loss为参考指标来更新学习率
#可视化显示训练过程,包括epoch,step,loss等
print('\nEpoch:' + str(epoch) + '/' + str(epochs))
print('step:' + str(step) + '/' + str(epoch_size) + '|| cls_loss:%.4f || lr:%.4f' % (cls_loss, lr))
if step % 10 == 0: #当step为10的倍数时,保存模型,这个参数可以根据自己需要改
checkpoint = {'model':model.state_dict(), #将一些参数以字典形式保存进checkpoint
'optimizer':optimizer,
'epoch':epoch}
torch.save(checkpoint, #保存
'./model/Eopch%d-Step%d-cls_loss%.4f.pth' % (epoch, step, class_loss))
在其中需要注意的地方就是Variable方法的使用,高版本依然可以用,但是不建议了,可以参照上面的代码自己去修改,再就是学习率的调整。
Pytorch中的学习率调整:lr_scheduler,ReduceLROnPlateau
torch.optim.lr_scheduler:该方法中提供了多种基于epoch训练次数进行学习率调整的方法;
torch.optim.lr_scheduler.ReduceLROnPlateau:该方法提供了一些基于训练过程中测量值对学习率进行动态的下降.
lr_scheduler调整方法一:根据epochs
CLASS torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1)
将每个参数组的学习率设置为给定函数的初始值,当last_epoch=-1时,设置初始的lr作为lr;
参数:
optimizer:封装好的优化器
lr_lambda(function or list):一个计算每个epoch的学习率的函数或者一个list;
last_epoch:最后一个epoch的索引
CLASS torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1)
当epoch每过stop_size时,学习率都变为初始学习率的gamma倍
CLASS torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma=0.1, last_epoch=-1)
当训练epoch达到milestones值时,初始学习率乘以gamma得到新的学习率;
CLASS torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma, last_epoch=-1)
每个epoch学习率都变为初始学习率的gamma倍
lr_scheduler调整方法二:根据测试指标
CLASS torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10,
verbose=False, threshold=0.0001, threshold_mode='rel',
cooldown=0, min_lr=0, eps=1e-08)
当参考的评价指标停止改进时,降低学习率,factor为每次下降的比例,训练过程中,当指标连续patience次数还没有改进时,降低学习率;
8.test.py 测试代码
这部分就是用于对模型进行实际测试的代码,将预测结果进行可视化操作。
from PIL import Image,ImageDraw,ImageFont
import torch
from vgg16_net import Vgg16_Net
import config as cfg
from torchvision.transforms import ToTensor
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #gpu可用使用gpu反之使用cpu
model_path = './model/Eopch1-Step1-cls_loss0.6872.pth' #模型的路径
model = Vgg16_Net() #定义model结构
checkpoint = torch.load(model_path) #加载checkpoint
model.load_state_dict(checkpoint['model']) #从checkpoint中加载进权重
net = model.eval() #测试时一定要加这一步,用于设置bn层不改变权值
net = net.to(device) #将模型防于gpu或者cpu上
to_tensor = ToTensor() #实例化ToTensor类
img_path = './data/Cat/Dog/1.jpg' #测试图片路径
img = Image.open(img_path) #打开图片
class_names = cfg.CLASS_NAMES #类别名
with torch.no_grad(): #不再进行反向传播
image = img.convert('RGB') #将图片转换为RGB格式
image = to_tensor(image) #转换为torch需要的tensor
image.to(device) #将图片防御gpu或者cpu上进行计算
outputs =net(image.unsqueeze(0)) #对图像进行扩维后进行反向传播
pred = torch.max(outputs, 1) #对预测值中取最大预测结果,即概率最大的那个值
pred = pred.indices.item() #取出编码后的类别
label = class_names[pred] #将编码的类别转换为原始的label
print(label)
font = ImageFont.truetype('arialuni.ttf', 36) #设置字体
draw = ImageDraw.Draw(img) #画图
draw.text((0,0), str(label), font=font, fill="red") #在图上显示出预测的label
img.show() #显示图片
下面是最后测试的可视化结果图:
以上就是整个二分类项目的所有代码,整个代码内容比较简单,主要是用来学习torch的代码结构和项目包括的内容,感受整个过程对于理解torch和学习torch很有帮助,希望大家看完整篇文章也能有一定的收获,当然大家也在学习torch的话,不妨也动动手仿造本案例手撕一下代码,等结束后你就会发现自己的收获,最后祝愿大家都能收获满满,成为深度学习领域的大牛,也欢迎大家在评论处询问或者提出自己的看法与建议。
放一个项目的传送门:提取码:0qip
更多技术欢迎加入交流:320297153