使用ONNX和Caffe2对PyTorch训练的模型进行推断应用

在2016年10月Face book发布PyTorch之后，由于其面向开发者友好，它很快获得了广泛应用。得益于良好的Python接口，它很适合用于研究和制作快速原型。在PyTorch中调试您的代码和测试网络模型架构可以非常容易地完成。 然而，当它投入生产时，谷歌的Tensorflow领先。使用TensorFlow服务部署机器学习模型非常容易。 这在2018年5月发生了变化，当时PyTorch与Caffe2集成，并获得了完整的生产流程。这是Face book使用的生产流程，他们使用PyTorch训练模型，并使用Caffe2部署模型。 注：Caffe2不应与Caffe混淆。它们是两个完全不同的框架。caffe在5年前很受欢迎，但现在它似乎已经失宠了。

1、Facebook的深度学习工作流

对于开发人员和工程师来说，需要在易于使用的语言中使用一个易于使用的开发框架来提高开发效率。PyTorch提供了这个解决方案。 然而，在生产部署中，Facebook需要以难以想象的规模运作。在这样的环境所追求的，是计算效率和性能，而不是开发人员的易用性。Facebook对性能问题的回答是Caffe2，这是一个非常快速高效的框架，用C++语言编写，用于生产。 Caffe2于2017年4月由Facebook推出。它是通用的，Caffe2模型可以部署在许多平台，包括移动端。Facebook在Caffe2的应用已经部署在超过10亿部iOS和Android手机上。 Face book保持PyTorch和Caffe2之间的互操作性。 在2019年5月，随着PyTorch1.1支持TensorBoard的发布，可视化和调试变得非常易用。

2、开放神经网络交换（ONN X）

开放神经网络交换（ONN X）是一种开放的格式，用户可以在不同的框架之间移动深度学习模型。这种开放格式最初是由Face book和微软提出的，但现在是一个被广泛接受的行业标准。 对于PyTorch模型的部署，最常见的方法是将它们转换为ONNX格式，然后使用Caffe2部署导出的ON NX模型。 在我们的上一篇文章中，我们描述了如何在PyTorch中训练图像分类器并进行推理。模型保存为.pt或.pth文件。在这篇文章中，我们将解释如何将经过训练的PyTorch模型转换为ONNX模型，并在Caffe2中进行推理。 我们还将使用PyTorch模型和ONNX模型来检验推理结果的异同。 我们使用PyTorch1.4.0+cpu 和onnx1.6.0和Python3.6进行这项工作。

3、PyTorch to ONNX

让我们看看如何将PyTorch.pt模型导出到ONNX。

 Export an ONNX model from a PyTorch .pt model

import torch.onnx
# Loading the input PyTorch model and mapping the tensors to CPU
device = torch.device('cpu')
model = torch.load('caltech_10_model_8.pt', map_location=device)

# Generate a dummy input that is consistent with the network's arhitecture
dummy_input = torch.randn(1, 3, 224, 224)

# Export into an ONNX model using the PyTorch model and the dummy input
torch.onnx.export(model, dummy_input, "animals_caltech.onnx",  export_params=True, keep_initializers_as_inputs=True)

由于我们将使用Caffe2在CPU中进行推理，我们将设备设置为“CPU”，并将Py Torch模型映射到CPU。然后，我们需要做一个虚拟输入，以适应网络结构的输入。最后，导出函数的参数包含PyTorch模型、虚拟输入和目标ON NX文件。keep_initializers_as_inputs=True 非常重要，不可缺少，否则，caffe2加载导出的模型时会报如下错误。

>>> rep = backend.prepare(model)
Traceback (most recent call last):
  File "", line 1, in 
  File "/usr/local/lib/python3.6/dist-packages/caffe2/python/onnx/backend.py", line 713, in prepare
    init_net, predict_net = cls._onnx_model_to_caffe2_net(model, device, opset_version, False)
  File "/usr/local/lib/python3.6/dist-packages/caffe2/python/onnx/backend.py", line 876, in _onnx_model_to_caffe2_net
    onnx_model = onnx.utils.polish_model(onnx_model)
  File "/usr/local/lib/python3.6/dist-packages/onnx/utils.py", line 21, in polish_model
    model = onnx.optimizer.optimize(model)
  File "/usr/local/lib/python3.6/dist-packages/onnx/optimizer.py", line 55, in optimize
    optimized_model_str = C.optimize(model_str, passes)
IndexError: _Map_base::at

请注意，对于安装Caffe2，目前预先构建的二进制文件没有CUDA对Mac、Ubuntu和CentOS的支持。所有其他平台或着需要CUDA支持都需要从源代码编译。在本例中，我们在Ubuntu18.04中测试了CPU推断。如果您想使用CUDA版本的ONNX模型，则需要从源构建Caffe2。

4、使用ONNX在Caffe2中的推断

接下来，我们现在可以在各种设备中部署我们的ON NX模型，并在Caffe2中进行推理。

# Inference in Caffe using the ONNX model

import caffe2.python.onnx.backend as backend
import onnx

from torchvision import transforms
from PIL import Image
import time
import numpy as np

# First load the onnx model
model = onnx.load("animals_caltech.onnx")

# Prepare the backend
rep = backend.prepare(model, device="CPU")

# Transform the image
transform = transforms.Compose([
        transforms.Resize(size=224),
        transforms.CenterCrop(size=224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
          
# Load and show the image
test_image_name = "caltech_10/test/zebra/250_0091.jpg"
test_image = Image.open(test_image_name)

# Apply the transformations to the input image and convert it into a tensor
test_image_tensor = transform(test_image)

# Make the input image ready to be input as a batch of size 1
test_image_tensor = test_image_tensor.view(1, 3, 224, 224)

time_start = time.time()
# Convert the tensor to numpy array
np_image = test_image_tensor.numpy()

# Pass the numpy array to run through the ONNX model
outputs = rep.run(np_image.astype(np.float32))
time_end = time.time()
print("Local CPU Inference time using ONNX model : ", time_end - time_start)

# Dictionary with class name and index
idx_to_class = {0: 'bear   ', 1: 'chimp  ', 2: 'giraffe', 3: 'gorilla', 4: 'llama  ', 5: 'ostrich', 6: 'porcupine', 7: 'skunk  ', 8: 'triceratops', 9: 'zebra  '}

ps = torch.exp(torch.from_numpy(outputs[0]))
topk, topclass = ps.topk(10, dim=1)
for i in range(10):
    print("Prediction", '{:2d}'.format(i+1), ":", '{:11}'.format(idx_to_class[topclass.cpu().numpy()[0][i]]), ", Class Id : ", topclass[0][i].numpy(), " Score: ", topk.cpu().detach().numpy()[0][i])

结果如下：

Local CPU Inference time using ONNX model :  0.25966334342956543
Prediction  1 : zebra       , Class Id :  9  Score:  0.9993488
Prediction  2 : triceratops , Class Id :  8  Score:  0.0002103801
Prediction  3 : giraffe     , Class Id :  2  Score:  0.0001857687
Prediction  4 : ostrich     , Class Id :  5  Score:  9.373466e-05
Prediction  5 : skunk       , Class Id :  7  Score:  4.9951355e-05
Prediction  6 : llama       , Class Id :  4  Score:  4.8880072e-05
Prediction  7 : chimp       , Class Id :  1  Score:  2.3821109e-05
Prediction  8 : porcupine   , Class Id :  6  Score:  1.8212046e-05
Prediction  9 : bear        , Class Id :  0  Score:  1.0552433e-05
Prediction 10 : gorilla     , Class Id :  3  Score:  9.918323e-06

 

>>> import caffe2.python.onnx.backend as backend
Traceback (most recent call last):
  File "", line 1, in 
  File "/usr/local/lib/python3.6/dist-packages/caffe2/python/onnx/backend.py", line 28, in 
    from caffe2.python import core, workspace, rnn_cell, gru_cell
  File "/usr/local/lib/python3.6/dist-packages/caffe2/python/core.py", line 9, in 
    from past.builtins import basestring
ModuleNotFoundError: No module named 'past'

如果出现如上错误，请使用

pip3 install future

我们加载ONNX模型并将其与设备信息一起传递给Caffe2。在本例中，它需要CPU，因为我们在生成ONNX模型时是为CPU导出的。 然后，我们可以读取输入测试图像，调整其大小，使其宽和高中较小的长度为224，在调整大小时保持纵横比。中心224×224图像被裁剪出来并转换成张量。此步骤将值转换为0-1的范围。然后，使用ImageNet平均值和标准差对图像进行正则化。按照输入[通道]=（输入[通道]-平均[通道]）/std[通道）进行计算。 然后使图像张量看起来像一批图像，因为网络模型要求输入一批图像。 然后将张量转换为Float32的numpy数组，并在Caffe2中运行加载的模型。 模型的输出是对数概率的形式。我们以他们的指数来获得实际的分数，对分数进行排序，并将分数最高的结果作为我们对输入测试图像的预测。 我们打印所有10个类的分数，按降序排列，以便我们可以直接使用PyTorch模型进行推理时计算的分数（就像我们以前的帖子一样），与使用Caffe2中的ONNX模型进行推理计算的分数进行比较。

5、PyTorch使用.pt模型进行推理

pytorch使用原模型进行推理如下：

import torch
from torchvision import transforms
from PIL import Image
import time
import numpy as np

# Loading the input PyTorch model and mapping the tensors to CPU
device = torch.device('cpu')
model = torch.load('caltech_10_model_8.pt', map_location=device)

# Transform the image
transform = transforms.Compose([
        transforms.Resize(size=224),
        transforms.CenterCrop(size=224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
          
# Load and show the image
test_image_name = "caltech_10/test/zebra/250_0091.jpg"
test_image = Image.open(test_image_name)
display(test_image)

# Apply the transformations to the input image and convert it into a tensor
test_image_tensor = transform(test_image)

# Make the input image ready to be input as a batch of size 1
test_image_tensor = test_image_tensor.view(1, 3, 224, 224)

# Dictionary with class name and index
idx_to_class = {0: 'bear   ', 1: 'chimp  ', 2: 'giraffe', 3: 'gorilla', 4: 'llama  ', 5: 'ostrich', 6: 'porcupine', 7: 'skunk  ', 8: 'triceratops', 9: 'zebra  '}

time_start = time.time()

with torch.no_grad():
    model.eval()
    # Model outputs log probabilities
    out = model(test_image_tensor)

time_end = time.time()

print("Local CPU Inference time  using Pytorch model : ", time_end - time_start)


ps = torch.exp(out)
topk, topclass = ps.topk(10, dim=1)
for i in range(10):
    print("Prediction", '{:2d}'.format(i+1), ":", '{:11}'.format(idx_to_class[topclass.cpu().numpy()[0][i]]), ", Class Id : ", topclass[0][i].numpy(), " Score: ", topk.cpu().detach().numpy()[0][i])

结果如下：

Local CPU Inference time  using Pytorch model :  0.37128472328186035

Prediction  1 : zebra       , Class Id :  9  Score:  0.9993488
Prediction  2 : triceratops , Class Id :  8  Score:  0.0002103801
Prediction  3 : giraffe     , Class Id :  2  Score:  0.00018576835
Prediction  4 : ostrich     , Class Id :  5  Score:  9.373448e-05
Prediction  5 : skunk       , Class Id :  7  Score:  4.9951497e-05
Prediction  6 : llama       , Class Id :  4  Score:  4.8880025e-05
Prediction  7 : chimp       , Class Id :  1  Score:  2.3821109e-05
Prediction  8 : porcupine   , Class Id :  6  Score:  1.8212062e-05
Prediction  9 : bear        , Class Id :  0  Score:  1.0552433e-05
Prediction 10 : gorilla     , Class Id :  3  Score:  9.918323e-06

6、结果对比

2种模型的准确度几乎完全一致，但caffe具有明显的速度优势。在Facebook和其他公司的移动和大规模应用程序中，ONNX模型已广泛部署在Caffe2运行时。

 

参考连接：https://www.learnopencv.com/pytorch-model-inference-using-onnx-and-caffe2/