如何將pytorch模型轉換爲tensorrt能夠掛載的模型

如何將pytorch模型轉換爲tensorrt能夠掛載的模型

github鏈接

安裝必要的包

1. 安裝pytorch/tensorflow/onnx/onnx_tf
2. python環境下，爲了便於安裝，可以將安裝鏡像改爲國內的鏡像，具體方法參見link
3. 安裝時，直接pip安裝即可，以下例子中所用到的版本tensorflow-gpu==1.15,onnx_tf==1.3

模型轉換

1. pytorch模型與tensorflow之間的轉換需要中間協議onnx，我們的轉換步驟爲：

   • pytorch模型轉換爲onnx模型
   • 需要將onnx模型的batch size，由固定轉爲可變
   • 將onnx模型轉換爲tensorflow模型，此時的tensorflow模型爲frozen的形式，當需要tensorflow serving或tensorrt掛載的時候，還需要進一步

2. pytorch模型轉換爲onnx模型

import os
import sys
sys.path.append('../')
from resnet import *
import torch

model = resnet34(num_classes=6, shortcut_type=True, sample_size=128, sample_duration=128)
weights = '../model/ct_pos_recogtion_20191115125435/ct_pos_recognition_0009_best.pth'
model.load_state_dict(torch.load(weights))

dummy_input = torch.randn(1,1,128,128,128)

torch.onnx.export(model, dummy_input, 'ctPosRecognition.onnx', verbose=True, input_names=['input'], output_names=['output'])

print('====> export to onnx model!')

3. onnx模型轉換爲pb

import onnx
from onnx_tf.backend import prepare

onnx_model = onnx.load('ctPosRecognition.onnx')
onnx_model.graph.input[0].type.tensor_type.shape.dim[0].dim_param = '?'
onnx.save(onnx_model, 'dynamic_model.onnx')

# 如果在jupyter或ipython中運行，需要通過CUDA_VISIBLE_DEVICES=0或os.environ['CUDA_VISIBLE_DEVICES']='0'，指定一個gpu
# 否則會報xla相關的錯誤
tf_rep = prepare(onnx_model)

print(tf_rep.inputs)
print(tf_rep.outputs)
print(tf_rep.tensor_dict)

print('\n====> inputs:')
for inp in tf_rep.inputs:
    print(tf_rep.tensor_dict[inp])

print('\n====> outputs:')
for outp in tf_rep.outputs:
    print(tf_rep.tensor_dict[outp])

tf_rep.export_graph('ctPosRecognition.pb')
print('export to frozen pb!')

4. pb文件轉換爲tensorflow serving或tensorrt掛載的文件

import os
import tensorflow as tf

infile = 'ctPosRecognition.pb'

graph = tf.get_default_graph()
sess = tf.Session()

with open(infile, 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    tf.import_graph_def(graph_def, name='')

inp = sess.graph.get_tensor_by_name('input:0')
print('====> input:\t')
print(inp)
outp = sess.graph.get_tensor_by_name('add_17:0')
print('====> output:\t')
print(outp)

export_path = './export'
builder = tf.saved_model.builder.SavedModelBuilder(export_path)

tensor_info_input = tf.saved_model.utils.build_tensor_info(inp)
tensor_info_output = tf.saved_model.utils.build_tensor_info(outp)

#定義簽名
prediction_signature = (
    tf.saved_model.signature_def_utils.build_signature_def(
    inputs={'images': tensor_info_input},
    outputs={'result': tensor_info_output},
    method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))

# 'serving_default' 可以隨意定義，但是直接利用系統的trt框架掛載，默認要填'serving_default'
builder.add_meta_graph_and_variables(sess, [tf.saved_model.tag_constants.SERVING],signature_def_map={'serving_default': prediction_signature})
builder.save()
print('Done exporting!')

檢查模型的輸出結果

過程中遇到的問題

1. pyFunc相關的問題， 主要時因爲pytorch做pooling的方式和tensorflow中pooling的VALID/SAME模型不同，一般在resnet等模型，第一層做pooling時，可以將pooling的padding改爲0
2. pytorch模型轉換爲onnx模型時，需要將固定的batch size轉換爲動態的batch size
3. 當使用多塊GPU卡進行onnx轉換爲tf的時候，可能會報xla的錯誤,可以通過環境變量的方式，將GPU限定到一塊卡，CUDA_VISIBLE_DEVICES=0或os.environ['CUDA_VISIBLE_DEVICES']='0'

reference

1. Dynamic dummy input when exporting a PyTorch model?

   • pytorch 轉 onnx時，batch_size 固定

2. 使用ONNX轉換PyTorch模型到 Tensorflow *.pb文件

3. Tensorflow的三種儲存格式-2（pb & Saved Model）