pytorch模型加载测试_「pytorch速成」Pytorch图像分类从模型自定义到测试

原标题&＃xff1a;「pytorch速成」Pytorch图像分类从模型自定义到测试

前面已跟大家介绍了Caffe和TensorFlow&＃xff0c;今天说说Pytorch。

1 什么是 Pytorch

一句话总结 Pytorch &＃61; Python &＃43; Torch。

Torch 是纽约大学的一个机器学习开源框架&＃xff0c;几年前在学术界非常流行&＃xff0c;包括 Lecun等大佬都在使用。但是由于使用的是一种绝大部分人绝对没有听过的 Lua 语言&＃xff0c;导致很多人都被吓退。后来随着 Python 的生态越来越完善&＃xff0c;Facebook 人工智能研究院推出了Pytorch并开源。Pytorch不是简单的封装 Torch并提供Python接口&＃xff0c;而是对Tensor以上的所有代码进行了重构&＃xff0c;同TensorFlow一样&＃xff0c;增加了自动求导。

后来Caffe2全部并入Pytorch&＃xff0c;如今已经成为了非常流行的框架。很多最新的研究如风格化、GAN 等大多数采用Pytorch源码&＃xff0c;这也是我们必须要讲解它的原因。

1.1 特点

(1) 动态图计算。TensorFlow从静态图发展到了动态图机制Eager Execution&＃xff0c;pytorch则一开始就是动态图机制。动态图机制的好处就是随时随地修改&＃xff0c;随处debug&＃xff0c;没有类似编译的过程。

(2) 简单。相比TensorFlow中Tensor、Variable、Session等概念充斥&＃xff0c;数据读取接口频繁更新&＃xff0c;tf.nn、tf.layers、tf.contrib各自重复&＃xff0c;Pytorch则是从Tensor到Variable再到nn.Module&＃xff0c;最新的Pytorch已经将Tensor和Variable合并&＃xff0c;这分别就是从数据张量到网络的抽象层次的递进。有人调侃TensorFlow的设计是“make it complicated”&＃xff0c;那么 Pytorch的设计就是“keep it simple”。

1.2 重要概念

(1) Tensor/Variable

每一个框架都有基本的数据结构&＃xff0c;Caffe是blob&＃xff0c;TensorFlow和Pytorch都是Tensor&＃xff0c;都是高维数组。Pytorch中的Tensor使用与Numpy的数组非常相似&＃xff0c;两者可以互转且共享内存。

tensor包括cpu和gpu两种类型&＃xff0c;如torch.FloatTensortorch.cuda.FloatTensorvirable&＃xff0c;就分别表示cpu和gpu下的32位浮点数。

tensor包含一些属性。data&＃xff0c;即Tensor内容&＃xff1b;Grad&＃xff0c;是与data对应的梯度&＃xff1b;requires_grad&＃xff0c;是否容许进行反向传播的学习&＃xff0c;更多的可以去查看API。

(2) nn.module

抽象好的网络数据结构&＃xff0c;可以表示为网络的一层&＃xff0c;也可以表示为一个网络结构&＃xff0c;这是一个基类。在实际使用过程中&＃xff0c;经常会定义自己的网络&＃xff0c;并继承nn.Module。具体的使用&＃xff0c;我们看下面的网络定义吧。

(3) torchvision包&＃xff0c;包含了目前流行的数据集&＃xff0c;模型结构和常用的图片转换工具

2 Pytorch 训练

安装咱们就不说了&＃xff0c;接下来的任务就是开始训练模型。训练模型包括数据准备、模型定义、结果保存与分析。

2.1 数据读取

前面已经介绍了Caffe和TensorFlow的数据读取&＃xff0c;两者的输入都是图片list&＃xff0c;但是读取操作过程差异非常大&＃xff0c;Pytorch与这两个又有很大的差异。这一次&＃xff0c;直接利用文件夹作为输入&＃xff0c;这是 Pytorch更加方便的做法。数据读取的完整代码如下&＃xff1a;

data_dir &＃61; &＃39;../../../../datas/head/&＃39;

data_transforms &＃61; {

&＃39;train&＃39;: transforms.Compose([

transforms.RandomSizedCrop(48),

transforms.RandomHorizontalFlip(),

transforms.ToTensor(),

transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])

]),

&＃39;val&＃39;: transforms.Compose([

transforms.Scale(64),

transforms.CenterCrop(48),

transforms.ToTensor(),

transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])

]),

}

image_datasets &＃61; {x: datasets.ImageFolder(os.path.join(data_dir, x),

data_transforms[x]) for x in [&＃39;train&＃39;, &＃39;val&＃39;]}

dataloders &＃61; {x: torch.utils.data.DataLoader(image_datasets[x],

batch_size&＃61;16,

shuffle&＃61;True,

num_workers&＃61;4) for x in [&＃39;train&＃39;, &＃39;val&＃39;]}

下面一个一个解释&＃xff0c;完整代码请移步 Git 工程。

(1) datasets.ImageFolder

Pytorch的torchvision模块中提供了一个dataset 包&＃xff0c;它包含了一些基本的数据集如mnist、coco、imagenet和一个通用的数据加载器ImageFolder。

它会以这样的形式组织数据&＃xff0c;具体的请到Git工程中查看。

root/left/1.png

root/left/2.png

root/left/3.png

root/right/1.png

root/right/2.png

root/right/3.png

imagefolder有3个成员变量。

self.classes&＃xff1a;用一个list保存类名&＃xff0c;就是文件夹的名字。

self.class_to_idx&＃xff1a;类名对应的索引&＃xff0c;可以理解为 0、1、2、3 等。

self.imgs&＃xff1a;保存(imgpath&＃xff0c;class)&＃xff0c;是图片和类别的数组。

不同文件夹下的图&＃xff0c;会被当作不同的类&＃xff0c;天生就用于图像分类任务。

(2) Transforms

这一点跟Caffe非常类似&＃xff0c;就是定义了一系列数据集的预处理和增强操作。到此&＃xff0c;数据接口就定义完毕了&＃xff0c;接下来在训练代码中看如何使用迭代器进行数据读取就可以了&＃xff0c;包括 scale、减均值等。

(3) torch.utils.data.DataLoader

这就是创建了一个 batch&＃xff0c;生成真正网络的输入。关于更多 Pytorch 的数据读取方法&＃xff0c;请自行学习。

2.2 模型定义

import torch

import torch.nn as nn

import torch.nn.functional as F

import numpy as np

class simpleconv3(nn.Module):&＃96;

def __init__(self):

super(simpleconv3,self).__init__()

self.conv1 &＃61; nn.Conv2d(3, 12, 3, 2)

self.bn1 &＃61; nn.BatchNorm2d(12)

self.conv2 &＃61; nn.Conv2d(12, 24, 3, 2)

self.bn2 &＃61; nn.BatchNorm2d(24)

self.conv3 &＃61; nn.Conv2d(24, 48, 3, 2)

self.bn3 &＃61; nn.BatchNorm2d(48)

self.fc1 &＃61; nn.Linear(48 * 5 * 5 , 1200)

self.fc2 &＃61; nn.Linear(1200 , 128)

self.fc3 &＃61; nn.Linear(128 , 2)

def forward(self , x):

x &＃61; F.relu(self.bn1(self.conv1(x)))

x &＃61; F.relu(self.bn1(self.conv2(x)))

x &＃61; F.relu(self.bn3(self.conv3(x)))

x &＃61; x.view(-1 , 48 * 5 * 5)

x &＃61; F.relu(self.fc1(x))

x &＃61; F.relu(self.fc2(x))

x &＃61; self.fc3(x)

return x

我们的例子都是采用一个简单的3层卷积 &＃43; 2层全连接层的网络结构。根据上面的网络结构的定义&＃xff0c;需要做以下事情。

(1) simpleconv3(nn.Module)

继承nn.Module&＃xff0c;前面已经说过&＃xff0c;Pytorch的网络层是包含在nn.Module 里&＃xff0c;所以所有的网络定义&＃xff0c;都需要继承该网络层&＃xff0c;并实现super方法&＃xff0c;如下&＃xff1a;

super(simpleconv3,self).__init__()

这个就当作一个标准执行就可以了。

(2) 网络结构的定义都在nn包里&＃xff0c;举例说明&＃xff1a;

torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride&＃61;1, padding&＃61;0, dilation&＃61;1, groups&＃61;1, bias&＃61;True)

完整的接口如上&＃xff0c;定义的第一个卷积层如下&＃xff1a;

nn.Conv2d(3, 12, 3, 2)

即输入通道为3&＃xff0c;输出通道为12&＃xff0c;卷积核大小为3&＃xff0c;stride&＃61;2&＃xff0c;其他的层就不一一介绍了&＃xff0c;大家可以自己去看nn的API。

(3) forward

backward方法不需要自己实现&＃xff0c;但是forward函数是必须要自己实现的&＃xff0c;从上面可以看出&＃xff0c;forward 函数也是非常简单&＃xff0c;串接各个网络层就可以了。

对比Caffe和TensorFlow可以看出&＃xff0c;Pytorch的网络定义更加简单&＃xff0c;初始化方法都没有显示出现&＃xff0c;因为 Pytorch已经提供了默认初始化。

如果我们想实现自己的初始化&＃xff0c;可以这么做&＃xff1a;

init.xavier_uniform(self.conv1.weight)init.constant(self.conv1.bias, 0.1)

它会对conv1的权重和偏置进行初始化。如果要对所有conv层使用 xavier 初始化呢&＃xff1f;可以定义一个函数&＃xff1a;

def weights_init(m):

if isinstance(m, nn.Conv2d):

xavier(m.weight.data)

xavier(m.bias.data)

net &＃61; Net()

net.apply(weights_init)

3 模型训练

网络定义和数据加载都定义好之后&＃xff0c;就可以进行训练了&＃xff0c;老规矩先上代码&＃xff1a;

def train_model(model, criterion, optimizer, scheduler, num_epochs&＃61;25):

for epoch in range(num_epochs):

print(&＃39;Epoch {}/{}&＃39;.format(epoch, num_epochs - 1))

for phase in [&＃39;train&＃39;, &＃39;val&＃39;]:

if phase &＃61;&＃61; &＃39;train&＃39;:

scheduler.step()

model.train(True)

else:

model.train(False)

running_loss &＃61; 0.0 running_corrects &＃61; 0.0

for data in dataloders[phase]:

inputs, labels &＃61; data

if use_gpu:

inputs &＃61; Variable(inputs.cuda())

labels &＃61; Variable(labels.cuda())

else:

inputs, labels &＃61; Variable(inputs), Variable(labels)

optimizer.zero_grad()

outputs &＃61; model(inputs)

_, preds &＃61; torch.max(outputs.data, 1)

loss &＃61; criterion(outputs, labels)

if phase &＃61;&＃61; &＃39;train&＃39;:

loss.backward()

optimizer.step()

running_loss &＃43;&＃61; loss.data.item()

running_corrects &＃43;&＃61; torch.sum(preds &＃61;&＃61; labels).item()

epoch_loss &＃61; running_loss / dataset_sizes[phase]

epoch_acc &＃61; running_corrects / dataset_sizes[phase]

if phase &＃61;&＃61; &＃39;train&＃39;:

writer.add_scalar(&＃39;data/trainloss&＃39;, epoch_loss, epoch)

writer.add_scalar(&＃39;data/trainacc&＃39;, epoch_acc, epoch)

else:

writer.add_scalar(&＃39;data/valloss&＃39;, epoch_loss, epoch)

writer.add_scalar(&＃39;data/valacc&＃39;, epoch_acc, epoch)

print(&＃39;{} Loss: {:.4f} Acc: {:.4f}&＃39;.format(

phase, epoch_loss, epoch_acc))

writer.export_scalars_to_json("./all_scalars.json")

writer.close()

return model

分析一下上面的代码&＃xff0c;外层循环是epoches&＃xff0c;然后利用 for data in dataloders[phase] 循环取一个epoch 的数据&＃xff0c;并塞入variable&＃xff0c;送入model。需要注意的是&＃xff0c;每一次forward要将梯度清零&＃xff0c;即optimizer.zero_grad()&＃xff0c;因为梯度会记录前一次的状态&＃xff0c;然后计算loss进行反向传播。

loss.backward()

optimizer.step()

下面可以分别得到预测结果和loss&＃xff0c;每一次epoch 完成计算。

epoch_loss &＃61; running_loss / dataset_sizes[phase]

epoch_acc &＃61; running_corrects / dataset_sizes[phase]

_, preds &＃61; torch.max(outputs.data, 1)

loss &＃61; criterion(outputs, labels)

可视化是非常重要的&＃xff0c;鉴于TensorFlow的可视化非常方便&＃xff0c;我们选择了一个开源工具包&＃xff0c;tensorboardx&＃xff0c;安装方法为pip install tensorboardx&＃xff0c;使用非常简单。

第一步&＃xff0c;引入包定义创建&＃xff1a;

from tensorboardX import SummaryWriter

writer &＃61; SummaryWriter()

第二步&＃xff0c;记录变量&＃xff0c;如train阶段的 loss&＃xff0c;writer.add_scalar(&＃39;data/trainloss&＃39;, epoch_loss, epoch)。

按照以上操作就完成了&＃xff0c;完整代码可以看配套的Git 项目&＃xff0c;我们看看训练中的记录。Loss和acc的曲线图如下&＃xff1a;

网络的收敛没有Caffe和TensorFlow好&＃xff0c;大家可以自己去调试调试参数了&＃xff0c;随便折腾吧。

4 Pytorch 测试

上面已经训练好了模型&＃xff0c;接下来的目标就是要用它来做inference了&＃xff0c;同样给出代码。

import torch

import torch.nn as nn

import torch.optim as optim

from torch.optim import lr_scheduler

from torch.autograd import Variable

import torchvision

from torchvision import datasets, models, transforms

import time

import os

from PIL import Image

import sys

import torch.nn.functional as F

from net import simpleconv3

data_transforms &＃61; transforms.Compose([

transforms.Resize(48),

transforms.ToTensor(),

transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])])

net &＃61; simpleconv3()

modelpath &＃61; sys.argv[1]

net.load_state_dict(torch.load(modelpath,map_location&＃61;lambda storage,loc: storage))

imagepath &＃61; sys.argv[2]

image &＃61; Image.open(imagepath)

imgblob &＃61; data_transforms(image).unsqueeze(0)

imgblob &＃61; Variable(imgblob)

torch.no_grad()

predict &＃61; F.softmax(net(imgblob))

print(predict)

从上面的代码可知&＃xff0c;做了几件事&＃xff1a;

定义网络并使用torch.load和load_state_dict载入模型。

用PIL的Image包读取图片&＃xff0c;这里没有用OpenCV&＃xff0c;因为Pytorch默认的图片读取工具就是PIL的Image&＃xff0c;它会将图片按照RGB的格式&＃xff0c;归一化到 0&＃xff5e;1 之间。读取图片之后&＃xff0c;必须转化为Tensor变量。

evaluation的时候&＃xff0c;必须设置torch.no_grad()&＃xff0c;然后就可以调用 softmax 函数得到结果了。

5 总结

本节讲了如何用 Pytorch 完成一个分类任务&＃xff0c;并学习了可视化以及使用训练好的模型做测试。

配套资料在github&＃xff0c;https://github.com/longpeng2008/yousan.ai。

责任编辑&＃xff1a;