python模型部署_pytorch模型部署到移动端（使用腾讯TNN）

记录分为

pytorch2onnx

onnx2tnn

tnn结果验证

移动端(安卓)使用

1、pytorch2onnx

环境&＃xff1a;

pytorch 1.4.0

onnx 1.6.0 (转换)

onnxruntime 1.3.0 (测试)

onnx-simplifier 0.2.9 (模型量化&＃xff0c;不执行后续报错了&＃xff0c;我测试是这样的)

转换代码&＃xff1a;

import onnx

import torch

from test_net import TestModel

import numpy as np

import cv2

if 1:

torch_model &＃61; TestModel("model.pt")

torch_model.eval()

batch_size &＃61; 1 #批处理大小

input_shape &＃61; (3,384,384) #输入数据

# set the model to inference mode

# torch_model.eval()

x &＃61; torch.randn(batch_size,*input_shape)# 生成张量

export_onnx_file &＃61; "./model.onnx"# 目的ONNX文件名

torch.onnx.export(torch_model,

x,

export_onnx_file,

export_params&＃61;True,

opset_version&＃61;11,

do_constant_folding&＃61;True, # wether to execute constant folding for optimization

input_names &＃61; [&＃39;input&＃39;], # the model&＃39;s input names

output_names &＃61; [&＃39;output&＃39;], # the model&＃39;s output names

dynamic_axes&＃61;{&＃39;input&＃39; : {0 : &＃39;batch_size&＃39;}, # variable lenght axes

&＃39;output&＃39; : {0 : &＃39;batch_size&＃39;}}

)

print (&＃39;get onnx ok!&＃39;)

利用 onnxruntime 测试转换的模型:

import onnxruntime

import imageio

import time

(width, height) &＃61; (384,384)

cap &＃61; cv2.VideoCapture(0)

while 1:

ret,img &＃61; cap.read()

time_start &＃61; time.time()

if img is None:

print(&＃39;no image input!&＃39;)

break

if img.ndim &＃61;&＃61; 2:

img &＃61; cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

in_height ,in_width ,_ &＃61; img.shape

img &＃61; cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0

img_resized &＃61; cv2.resize(img, (width, height), interpolation&＃61;cv2.INTER_AREA)

img_resized &＃61; (

torch.from_numpy(np.transpose(img_resized, (2, 0, 1))).contiguous().float()

)

value &＃61; img_resized.unsqueeze(0)

def to_numpy(tensor):

return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

ort_session &＃61; onnxruntime.InferenceSession("model.onnx")

ort_inputs &＃61; {ort_session.get_inputs()[0].name: (to_numpy(value)).astype(np.float32)}

#Actual: (N11onnxruntime17PrimitiveDataTypeIdEE) , expected: (N11onnxruntime17PrimitiveDataTypeIfEE)

#传入数据类型不对

ort_outs &＃61; ort_session.run(None, ort_inputs)

result &＃61; ort_outs[0][0, 0, :, :]

result &＃61; np.array(result)

print (result.shape)

reslut_resized &＃61; cv2.resize(

result, (in_width, in_height), interpolation&＃61;cv2.INTER_AREA

)

print(&＃39;cost : %.3f (s)&＃39;%(time.time() - time_start))

cv2.namedWindow(&＃39;re&＃39;,2)

cv2.imshow(&＃39;re&＃39;,reslut_resized)

if cv2.waitKey(1) &＃61;&＃61;27:

break

cap.release()

cv2.destroyAllWindows()

模型简化操作&＃xff1a;

python -m onnxsim model10.onnx model_sim.onnx --input-shape 1,3,384,384

2、onnx2tnn

下载tnn源码&＃xff0c;https://github.com/Tencent/TNN &＃xff1b;进入

~/TNN-master/tools/onnx2tnn/onnx-converter 文件夹&＃xff0c;运行 ./build 进行编译。

2.运行命令进行转换

python onnx2tnn.py model/model_sim.onnx -version&＃61;algo_version -optimize&＃61;1

0.----onnx version:1.6.0

结果为&＃xff1a;

algo_optimize 1

onnx_net_opt_path /home/jiang/TNN-master/tools/onnx2tnn/onnx-converter/model/model_sim.opt.onnx

1.----onnx_optimizer: /home/jiang/TNN-master/tools/onnx2tnn/onnx-converter/model/model_sim.onnx

/home/jiang/TNN-master/tools/onnx2tnn/onnx-converter

----load onnx model: /home/jiang/TNN-master/tools/onnx2tnn/onnx-converter/model/model_sim.onnx

----onnxsim.simplify error: You&＃39;d better check the result with Netron

----onnxsim.simplify error:

----export optimized onnx model: /home/jiang/TNN-master/tools/onnx2tnn/onnx-converter/model/model_sim.opt.onnx

----export optimized onnx model done

2.----onnx2tnn: /home/jiang/TNN-master/tools/onnx2tnn/onnx-converter/model/model_sim.opt.onnx

get_node_attr_ai [Line 116] name :546

get_node_attr_ai [Line 116] name :585

get_node_attr_ai [Line 116] name :624

get_node_attr_ai [Line 116] name :663

get_node_attr_ai [Line 116] name :693

TNNLayerParam [Line 61] resize: coordinate_transformation_mode(pytorch_half_pixel) is not supported, result may be different.

3.----onnx2tnn status: 0

出现了错误----onnxsim.simplify error: You’d better check the result with Netron

----onnxsim.simplify error: 。

参数 algo_optimize&＃61;0即不进行优化就不保存转换成功。

3、tnn结果验证

先要安装Android环境&＃xff1b;

本文地址&＃xff1a;https://blog.csdn.net/qq_32768679/article/details/107106390

如您对本文有疑问或者有任何想说的&＃xff0c;请点击进行留言回复&＃xff0c;万千网友为您解惑&＃xff01;