pytorch——训练第一个分离器讲解_PyTorch专栏（十三）：使用ONNX将模型转移至Caffe2和移动端...
作者 | News
编辑 | 奇予纪
出品 | 磐创AI团队出品
【磐创AI 导读】：本篇文章讲解了PyTorch专栏的第四章中的使用ONNX将模型转移至Caffe2和移动端。查看专栏历史文章，请点击下方蓝色字体进入相应链接阅读。查看关于本专栏的介绍：PyTorch专栏开篇。想要更多电子杂志的机器学习，深度学习资源，大家欢迎点击上方蓝字关注我们的公众号：磐创AI。
专栏目录： 第一章： PyTorch之简介与下载- PyTorch简介
- PyTorch环境搭建
- PyTorch入门
- PyTorch自动微分
- PyTorch神经网络
- PyTorch图像分类器
- PyTorch数据并行处理
- 数据加载和处理
- PyTorch小试牛刀
- 迁移学习
- 混合前端的seq2seq模型部署
- 保存和加载模型
- 微调基于torchvision 0.3的目标检测模型
- 微调TorchVision模型
- 空间变换器网络
- 使用PyTorch进行神经传递
- 生成对抗示例
- 使用ONNX将模型转移至Caffe2和移动端
- 聊天机器人教程
- 使用字符级RNN生成名字
- 使用字符级RNN进行名字分类
- 在深度学习和NLP中使用Pytorch
- 使用Sequence2Sequence网络和注意力进行翻译
使用ONNX将模型转移至Caffe2和移动端
在本教程中，我们将介绍如何使用 ONNX 将 PyTorch 中定义的模型转换为 ONNX 格式，然后将其加载到 Caffe2 中。 一旦进入 Caffe2，我们就可以运行模型来仔细检查它是否正确导出，然后我们展示了如何使用 Caffe2 功能(如移动导出器)在移动设备上执行模型。在本教程中，您需要安装onnx和Caffe2。您可以使用pip install onnx来获取 onnx。
注意：本教程需要 PyTorch master 分支，可以按照这里说明进行安装。
1.引入模型
# 一些包的导入
import io
import numpy as np
from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx
1.1 `SuperResolution`模型
超分辨率是一种提高图像、视频分辨率的方法，广泛用于图像处理或视频剪辑。在本教程中，我们将首先使用带有虚拟输入的小型超分辨率模型。
首先，让我们在 PyTorch 中创建一个SuperResolution模型。这个模型直接来自 PyTorch 的例子，没有修改：
# PyTorch中定义的Super Resolution模型
import torch.nn as nn
import torch.nn.init as init
class SuperResolutionNet(nn.Module):
def __init__(self, upscale_factor, inplace=False):
super(SuperResolutionNet, self).__init__()
self.relu = nn.ReLU(inplace=inplace)
self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
self._initialize_weights()
def forward(self, x):
x = self.relu(self.conv1(x))
x = self.relu(self.conv2(x))
x = self.relu(self.conv3(x))
x = self.pixel_shuffle(self.conv4(x))
return x
def _initialize_weights(self):
init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
init.orthogonal_(self.conv4.weight)
# 使用上面模型定义，创建super-resolution模型
torch_model = SuperResolutionNet(upscale_factor=3)
1.2 训练模型
通常，你现在会训练这个模型; 但是，对于本教程我们将下载一些预先训练的权重。请注意，此模型未经过充分训练来获得良好的准确性，此处仅用于演示目的。
# 加载预先训练好的模型权重
del_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
batch_size = 1 # just a random number
# 使用预训练的权重初始化模型
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
map_location = None
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))
# 将训练模式设置为falsesince we will only run the forward pass.
torch_model.train(False)
1.3 导出模型
在 PyTorch 中通过跟踪工作导出模型。要导出模型，请调用torch.onnx._export()函数。这将执行模型，记录运算符用于计算输出的轨迹。因为_export运行模型，我们需要提供输入张量x。这个张量的值并不重要; 它可以是图像或随机张量，只要它大小是正确的。
要了解有关PyTorch导出界面的更多详细信息，请查看torch.onnx documentation文档。
# 输入模型
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
# 导出模型
torch_out = torch.onnx._export(torch_model, # model being run
x, # model input (or a tuple for multiple inputs)
"super_resolution.onnx", # where to save the model (can be a file or file-like object)
export_params=True) # store the trained parameter weights inside the model file
torch_out是执行模型后的输出。通常您可以忽略此输出，但在这里我们将使用它来验证我们导出的模型在Caffe2中运行时是否计算出相同的值。
1.4 采用ONNX表示模型并在Caffe2中使用
现在让我们采用 ONNX 表示并在 Caffe2 中使用它。这部分通常可以在一个单独的进程中或在另一台机器上完成，但我们将在同一个进程中继续，以便我们可以验证 Caffe2 和 PyTorch 是否为网络计算出相同的值：
import onnx
import caffe2.python.onnx.backend as onnx_caffe2_backend
#加载ONNX ModelProto对象。模型是一个标准的Python protobuf对象
model = onnx.load("super_resolution.onnx")
# 为执行模型准备caffe2后端，将ONNX模型转换为可以执行它的Caffe2 NetDef。
# 其他ONNX后端，如CNTK的后端即将推出。
prepared_backend = onnx_caffe2_backend.prepare(model)
# 在Caffe2中运行模型
# 构造从输入名称到Tensor数据的映射。
# 模型图形本身包含输入图像之后所有权重参数的输入。由于权重已经嵌入，我们只需要传递输入图像。
# 设置第一个输入。
W = {model.graph.input[0].name: x.data.numpy()}
# 运行Caffe2 net:
c2_out = prepared_backend.run(W)[0]
# 验证数字正确性，最多3位小数
np.testing.assert_almost_equal(torch_out.data.cpu().numpy(), c2_out, decimal=3)
print("Exported model has been executed on Caffe2 backend, and the result looks good!")
我们应该看到 PyTorch 和 Caffe2 的输出在数字上匹配最多3位小数。作为旁注，如果它们不匹配则存在 Caffe2 和 PyTorch 中的运算符以不同方式实现的问题，请在这种情况下与我们联系。
2.使用ONNX转换SRResNET
使用与上述相同的过程，我们参考文章中提出的超分辨率转移了一个有趣的新模型“SRResNet”(感谢Twitter上的作者为本教程的目的提供了代码和预训练参数)。可在此处找到模型定义和预训练模型。下面是 SRResNet 模型的输入、输出。
3.在移动设备上运行模型
到目前为止，我们已经从 PyTorch 导出了一个模型，并展示了如何加载它并在 Caffe2 中运行它。现在模型已加载到 Caffe2 中，我们可以将其转换为适合在移动设备上运行的格式。
我们将使用 Caffe2 的mobile_exporter生成可在移动设备上运行的两个模型protobufs。第一个用于使用正确的权重初始化网络，第二个实际运行执行模型。在本教程的其余部分，我们将继续使用小型超分辨率模型。
# 从内部表示中提取工作空间和模型原型
c2_workspace = prepared_backend.workspace
c2_model = prepared_backend.predict_net
# 现在导入caffe2的`mobile_exporter`
from caffe2.python.predictor import mobile_exporter
# 调用Export来获取predict_net，init_net。 在移动设备上运行时需要这些网络
init_net, predict_net = mobile_exporter.Export(c2_workspace, c2_model, c2_model.external_input)
# 我们还将init_net和predict_net保存到我们稍后将用于在移动设备上运行它们的文件中
with open('init_net.pb', "wb") as fopen:
fopen.write(init_net.SerializeToString())
with open('predict_net.pb', "wb") as fopen:
fopen.write(predict_net.SerializeToString())
init_net具有模型参数和嵌入在其中的模型输入，predict_net将用于指导运行时的init_net执行。在本教程中，我们将使用上面生成的init_net和predict_net，并在正常的 Caffe2 后端和移动设备中运行它们，并验证两次运行中生成的输出高分辨率猫咪图像是否相同。
在本教程中，我们将使用广泛使用的著名猫咪图像，如下所示：
# 一些必备的导入包
from caffe2.proto import caffe2_pb2
from caffe2.python import core, net_drawer, net_printer, visualize, workspace, utils
import numpy as np
import os
import subprocess
from PIL import Image
from matplotlib import pyplot
from skimage import io, transform
3.1 加载图像并预处理
首先，让我们加载图像，使用标准的skimage python库对其进行预处理。请注意，此预处理是处理用于训练/测试神经网络的数据的标准做法。
# 加载图像
img_in = io.imread("./_static/img/cat.jpg")
# 设置图片分辨率为 224x224
img = transform.resize(img_in, [224, 224])
# 保存好设置的图片作为模型的输入
io.imsave("./_static/img/cat_224x224.jpg", img)
3.2 在Caffe2运行并输出
现在，作为下一步，让我们拍摄调整大小的猫图像并在 Caffe2 后端运行超分辨率模型并保存输出图像。这里的图像处理步骤已经从 PyTorch 实现的超分辨率模型中采用。
# 加载设置好的图片并更改为YCbCr的格式
img = Image.open("./_static/img/cat_224x224.jpg")
img_ycbcr = img.convert('YCbCr')
img_y, img_cb, img_cr = img_ycbcr.split()
# 让我们运行上面生成的移动网络，以便正确初始化caffe2工作区
workspace.RunNetOnce(init_net)
workspace.RunNetOnce(predict_net)
# Caffe2有一个很好的net_printer能够检查网络的外观
# 并确定我们的输入和输出blob名称是什么。
print(net_printer.to_string(predict_net))
从上面的输出中，我们可以看到输入名为“9”，输出名为“27”(我们将数字作为blob名称有点奇怪，但这是因为跟踪JIT为模型生成了编号条目)。
# 现在，让我们传递调整大小的猫图像以供模型处理。
workspace.FeedBlob("9", np.array(img_y)[np.newaxis, np.newaxis, :, :].astype(np.float32))
# 运行predict_net以获取模型输出
workspace.RunNetOnce(predict_net)
# 现在让我们得到模型输出blob
img_out = workspace.FetchBlob("27")
现在，我们将在这里回顾PyTorch实现超分辨率模型的后处理步骤，以构建最终输出图像并保存图像。
img_out_y = Image.fromarray(np.uint8((img_out[0, 0]).clip(0, 255)), mode='L')
# 获取输出图像遵循PyTorch实现的后处理步骤
final_img = Image.merge(
"YCbCr", [
img_out_y,
img_cb.resize(img_out_y.size, Image.BICUBIC),
img_cr.resize(img_out_y.size, Image.BICUBIC),
]).convert("RGB")
# 保存图像，我们将其与移动设备的输出图像进行比较
final_img.save("./_static/img/cat_superres.jpg")
3.3 在移动端上执行模型
我们已经完成了在纯Caffe2后端运行我们的移动网络，现在，让我们在Android设备上执行该模型并获取模型输出。
注意：对于 Android 开发，需要
adb shell，否则教程的以下部分将无法运行。
在我们在移动设备上运行模型的第一步中，我们把基于移动设备的本机速度测试基准二进制文件推送到 adb 。
这个二进制文件可以在移动设备上执行模型，也可以导出我们稍后可以检索的模型输出。二进制文件可在此处获得。要构建二进制文件，请按照此处的说明执行build_android.sh脚本。
注意：你需要已经安装了
ANDROID_NDK,并且设置环境变量ANDROID_NDK=path to ndk root。
# 让我们先把一堆东西推到adb，指定二进制的路径
CAFFE2_MOBILE_BINARY = ('caffe2/binaries/speed_benchmark')
# 我们已经在上面的步骤中保存了`init_net`和`proto_net`，我们现在使用它们。
# 推送二进制文件和模型protos
os.system('adb push ' + CAFFE2_MOBILE_BINARY + ' /data/local/tmp/')
os.system('adb push init_net.pb /data/local/tmp')
os.system('adb push predict_net.pb /data/local/tmp')
# 让我们将输入图像blob序列化为blob proto，然后将其发送到移动设备以供执行。
with open("input.blobproto", "wb") as fid:
fid.write(workspace.SerializeBlob("9"))
# 将输入图像blob推送到adb
os.system('adb push input.blobproto /data/local/tmp/')
# 现在我们在移动设备上运行网络，查看`speed_benchmark --help`，了解各种选项的含义
os.system(
'adb shell /data/local/tmp/speed_benchmark ' # binary to execute
'--init_net=/data/local/tmp/super_resolution_mobile_init.pb ' # mobile init_net
'--net=/data/local/tmp/super_resolution_mobile_predict.pb ' # mobile predict_net
'--input=9 ' # name of our input image blob
'--input_file=/data/local/tmp/input.blobproto ' # serialized input image
'--output_folder=/data/local/tmp ' # destination folder for saving mobile output
'--output=27,9 ' # output blobs we are interested in
'--iter=1 ' # number of net iterations to execute
'--caffe2_log_level=0 '
)
# 从adb获取模型输出并保存到文件
os.system('adb pull /data/local/tmp/27 ./output.blobproto')
# 我们可以使用与之前相同的步骤恢复输出内容并对模型进行后处理
blob_proto = caffe2_pb2.BlobProto()
blob_proto.ParseFromString(open('./output.blobproto').read())
img_out = utils.Caffe2TensorToNumpyArray(blob_proto.tensor)
img_out_y = Image.fromarray(np.uint8((img_out[0,0]).clip(0, 255)), mode='L')
final_img = Image.merge(
"YCbCr", [
img_out_y,
img_cb.resize(img_out_y.size, Image.BICUBIC),
img_cr.resize(img_out_y.size, Image.BICUBIC),
]).convert("RGB")
final_img.save("./_static/img/cat_superres_mobile.jpg")
现在，您可以比较图像 cat_superres.jpg(来自纯caffe2后端执行的模型输出)和 cat_superres_mobile.jpg(来自移动执行的模型输出)，并看到两个图像看起来相同。
如果它们看起来不一样，那么在移动设备上执行会出现问题，在这种情况下，请联系Caffe2社区。你应该期望看
使用上述步骤，您可以轻松地在移动设备上部署模型。另外，有关caffe2移动后端的更多信息，请查看caffe2-android-demo。
相关链接：
SuperResolution模型：https://github.com/pytorch/examples/blob/master/super_resolution/model.py
torch.onnx documentation文档：
https://pytorch.org/docs/master/onnx.html
“SRResNet”参考文章：
https://arxiv.org/pdf/1609.04802.pdf
“SRResNet”模型定义和预训练模型：https://gist.github.com/prigoyal/b245776903efbac00ee89699e001c9bd
PyTorch实现超分辨率模型的后处理步骤：https://github.com/pytorch/examples/blob/master/super_resolution/super_resolve.py
二进制文件：https://github.com/caffe2/caffe2/blob/master/caffe2/binaries/speed_benchmark.cc
二进制文件说明：https://github.com/caffe2/caffe2/blob/master/scripts/build_android.sh
caffe2-android-demo：
https://caffe2.ai/docs/AI-Camera-demo-android.html
下方点击 | 阅读原文 | 了解更多