python模型部署_pytorch模型部署到移动端(使用腾讯TNN)

记录分为

pytorch2onnx

onnx2tnn

tnn结果验证

移动端(安卓)使用

1、pytorch2onnx

环境:

pytorch 1.4.0

onnx 1.6.0 (转换)

onnxruntime 1.3.0 (测试)

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

转换代码:

import onnx

import torch

from test_net import TestModel

import numpy as np

import cv2

if 1:

torch_model = TestModel("model.pt")

torch_model.eval()

batch_size = 1 #批处理大小

input_shape = (3,384,384) #输入数据

# set the model to inference mode

# torch_model.eval()

x = torch.randn(batch_size,*input_shape)# 生成张量

export_onnx_file = "./model.onnx"# 目的ONNX文件名

torch.onnx.export(torch_model,

x,

export_onnx_file,

export_params=True,

opset_version=11,

do_constant_folding=True, # wether to execute constant folding for optimization

input_names = ['input'], # the model's input names

output_names = ['output'], # the model's output names

dynamic_axes={'input' : {0 : 'batch_size'}, # variable lenght axes

'output' : {0 : 'batch_size'}}

)

print ('get onnx ok!')

利用 onnxruntime 测试转换的模型:

import onnxruntime

import imageio

import time

(width, height) = (384,384)

cap = cv2.VideoCapture(0)

while 1:

ret,img = cap.read()

time_start = time.time()

if img is None:

print('no image input!')

break

if img.ndim == 2:

img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

in_height ,in_width ,_ = img.shape

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0

img_resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)

img_resized = (

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

)

value = img_resized.unsqueeze(0)

def to_numpy(tensor):

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

ort_session = onnxruntime.InferenceSession("model.onnx")

ort_inputs = {ort_session.get_inputs()[0].name: (to_numpy(value)).astype(np.float32)}

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

#传入数据类型不对

ort_outs = ort_session.run(None, ort_inputs)

result = ort_outs[0][0, 0, :, :]

result = np.array(result)

print (result.shape)

reslut_resized = cv2.resize(

result, (in_width, in_height), interpolation=cv2.INTER_AREA

)

print('cost : %.3f (s)'%(time.time() - time_start))

cv2.namedWindow('re',2)

cv2.imshow('re',reslut_resized)

if cv2.waitKey(1) ==27:

break

cap.release()

cv2.destroyAllWindows()

模型简化操作:

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

2、onnx2tnn

下载tnn源码,https://github.com/Tencent/TNN ;进入

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

2.运行命令进行转换

python onnx2tnn.py model/model_sim.onnx -version=algo_version -optimize=1

0.----onnx version:1.6.0

结果为:

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'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=0即不进行优化就不保存转换成功。

python模型部署_pytorch模型部署到移动端(使用腾讯TNN)_第1张图片

3、tnn结果验证

先要安装Android环境;

本文地址:https://blog.csdn.net/qq_32768679/article/details/107106390

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

你可能感兴趣的:(python模型部署)