使用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 xdef _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。

總結

以上是生活随笔為你收集整理的使用ONNX将模型转移至Caffe2和移动端的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。