我們訓練好深度學習模型後,這時其仍然需在特定的深度學習框架下運行,往往不能進行高性能推理。
NVIDIA提供了一套高效推理的框架——TensorRT,可將已訓練好的模型轉爲TensorRT引擎格式,然後進行高效推理。
對於Pytorch用戶而言,該技術路線爲:pytorch model-->onnx file-->TensorRT engine。
因此,我們需要做的只有三步:
- 將Pytorch模型轉爲ONNX作爲中間格式;
- 將ONNX文件轉爲TensorRT引擎(格式包括:FP32、FP16、INT8);
- 使用TensorRT引擎文件進行推理計算。
關於TensorRT的介紹網上資料較多,這裏就不再贅述。下面將結合這三個步驟對整個過程進行簡單介紹 。
詳細的代碼文件我已整理到GitHub(https://github.com/qq995431104/Pytorch2TensorRT.git),歡迎大家參考並給個star~~
目錄
1、Pytorch to ONNX
這一步比較簡單,只要你的模型中所有OP均被ONNX支持,即可利用Pytorch中的ONN庫進行轉換。參考如下代碼:
import torch
def get_model():
""" Define your own model and return it
:return: Your own model
"""
pass
def get_onnx(model, onnx_save_path, example_tensor):
example_tensor = example_tensor.cuda()
_ = torch.onnx.export(model, # model being run
example_tensor, # model input (or a tuple for multiple inputs)
onnx_save_path,
verbose=False, # store the trained parameter weights inside the model file
training=False,
do_constant_folding=True,
input_names=['input'],
output_names=['output']
)
if __name__ == '__main__':
model = get_model()
onnx_save_path = "onnx/resnet50_2.onnx"
example_tensor = torch.randn(1, 3, 288, 512, device='cuda')
# 導出模型
get_onnx(model, onnx_save_path, example_tensor)
需要提供的有:加載好的Pytorch模型、一個輸入樣例。其中,模型需要按照自己的方式導入並加載模型,輸入樣例的格式爲BCHW,B爲batch_size,CHW爲通道、高、寬,CHW的值需要與你自己的模型相匹配,否則後面轉換成功後輸出結果也不對。
如果出現“RuntimeError: ONNX export failed: Couldn't export Python operator XXXX”錯誤提示,說明你的模型中有ONNX不支持的OP,可以嘗試升級Pytorch版本,或者編寫自定義op,這部分尚未進行研究,後續有了進展會更新上來。
2、ONNX to TensorRT
現在已經有了ONNX文件了,需要利用TensorRT提供的OnnxParser解析該文件,同理:Caffe對應的有CaffPaser、TensorFlow的UFF格式對應的有UffParser。
先使用TensorRT創建一個builder,然後創建一個network,然後利用對應的Parser將ONNX文件加載進去;
接着,對builder指定一些參數設置,如:max_batch_size、max_workspace_size;
如需轉爲特定格式,如fp16或int8,需指定相應參數:fp16_mode或int8_mode設爲True;
對於Int8格式,需要:
- 準備一個校準集,用於在轉換過程中尋找使得轉換後的激活值分佈與原來的FP32類型的激活值分佈差異最小的閾值;
- 並寫一個校準器類,該類需繼承trt.IInt8EntropyCalibrator2父類,並重寫get_batch_size, get_batch, read_calibration_cache, write_calibration_cache這幾個方法。具體做法參考腳本
myCalibrator.py
. - 使用時,需額外指定cache_file,該參數是校準集cache文件的路徑,會在校準過程中生成,方便下一次校準時快速提取。
示例代碼如下:
def ONNX2TRT(args, calib=None):
''' convert onnx to tensorrt engine, use mode of ['fp32', 'fp16', 'int8']
:return: trt engine
'''
assert args.mode.lower() in ['fp32', 'fp16', 'int8'], "mode should be in ['fp32', 'fp16', 'int8']"
G_LOGGER = trt.Logger(trt.Logger.WARNING)
with trt.Builder(G_LOGGER) as builder, builder.create_network() as network, \
trt.OnnxParser(network, G_LOGGER) as parser:
builder.max_batch_size = args.batch_size
builder.max_workspace_size = 1 << 30
if args.mode.lower() == 'int8':
assert (builder.platform_has_fast_int8 == True), "not support int8"
builder.int8_mode = True
builder.int8_calibrator = calib
elif args.mode.lower() == 'fp16':
assert (builder.platform_has_fast_fp16 == True), "not support fp16"
builder.fp16_mode = True
print('Loading ONNX file from path {}...'.format(args.onnx_file_path))
with open(args.onnx_file_path, 'rb') as model:
print('Beginning ONNX file parsing')
parser.parse(model.read())
print('Completed parsing of ONNX file')
print('Building an engine from file {}; this may take a while...'.format(args.onnx_file_path))
engine = builder.build_cuda_engine(network)
print("Created engine success! ")
# 保存計劃文件
print('Saving TRT engine file to path {}...'.format(args.engine_file_path))
with open(args.engine_file_path, "wb") as f:
f.write(engine.serialize())
print('Engine file has already saved to {}!'.format(args.engine_file_path))
return engine
3、推理
推理過程就完全獨立於我們原先模型所依賴的框架了。
基本過程如下:
- 按照原模型的輸入輸出格式,準備數據,如:輸入的shape、均值、方差,輸出的shape等;
- 根據第二步得到的引擎文件,利用TensorRT Runtime反序列化爲引擎engine;
- 創建上下文環境engine.create_execution_context();
- 使用Pycuda的mem_alloc對輸入輸出分配cuda內存;
- 創建Stream;
- 使用memcpy_htod_async將IO數據放入device(一般爲GPU);
- 使用context.execute_async執行推理;
- 使用memcpy_dtoh_async取出結果;
根據引擎文件反序列化爲TensorRT引擎的示例代碼如下:
def loadEngine2TensorRT(filepath):
G_LOGGER = trt.Logger(trt.Logger.WARNING)
# 反序列化引擎
with open(filepath, "rb") as f, trt.Runtime(G_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(f.read())
return engine
推理過程示例如下:
# 通過engine文件創建引擎
engine = loadEngine2TensorRT('path_to_engine_file')
# 準備輸入輸出數據
img = Image.open('XXX.jpg')
img = D.transform(img).unsqueeze(0)
img = img.numpy()
output = np.empty((1, 2), dtype=np.float32)
#創建上下文
context = engine.create_execution_context()
# 分配內存
d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize)
d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)
bindings = [int(d_input), int(d_output)]
# pycuda操作緩衝區
stream = cuda.Stream()
# 將輸入數據放入device
cuda.memcpy_htod_async(d_input, img, stream)
# 執行模型
context.execute_async(100, bindings, stream.handle, None)
# 將預測結果從從緩衝區取出
cuda.memcpy_dtoh_async(output, d_output, stream)
# 線程同步
stream.synchronize()
print(output)
*更多詳細內容,請參閱GitHub倉庫:https://github.com/qq995431104/Pytorch2TensorRT.git