( Tencent -TNN 學習)pytorch模型部署到移動端

記錄分爲

  • 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官方文檔

  1. 下載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: <class 'RuntimeError'>
----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: <class ‘RuntimeError’>。
參數 algo_optimize=0即不進行優化就不保存轉換成功。
在這裏插入圖片描述

3、tnn結果驗證

1)、TNN編譯,不同平臺庫編譯見鏈接。以下爲安卓庫編譯:
環境要求

  • 依賴庫
    cmake(使用3.6及以上版本)
sudo apt-get install attr
  • NDK配置

    下載ndk版本(>=15c) https://developer.android.com/ndk/downloads
    配置環境變量 :

 sudo gedit ~/.bashrc

#ndk
export ANDROID_NDK=/home/jiang/Android/android-ndk-r21b
export PATH=PATH:{PATH}:ANDROID_NDK
#tnn
export TNN_ROOT_PATH=/home/jiang/TNN-master export
PATH=PATH:{PATH}:TNN_ROOT_PATH

  source ~/.bashrc
  • 編譯

切換到~/TNN-master/scripts下,修改腳本 build_android.sh;然後再執行./build_android.sh進行編譯。

ABIA32=“armeabi-v7a with NEON”
ABIA64=“arm64-v8a”
STL=“c++_static”
SHARED_LIB=“ON” # ON表示編譯動態庫,OFF表示編譯靜態庫
ARM=“ON” # ON表示編譯帶有Arm CPU版本的庫
OPENMP=“ON” # ON表示打開OpenMP
OPENCL=“ON” # ON表示編譯帶有Arm GPU版本的庫
SHARING_MEM_WITH_OPENGL=0 # 1表示OpenGL的Texture可以與OpenCL共享

編譯完成後,在當前目錄的release目錄下生成對應的armeabi-v7a庫,arm64-v8a庫和include頭文件。
在這裏插入圖片描述
2)tnn模型驗證

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章