Pytorch-猫狗分类实战（下）

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