TensorRT
TensorRT是NVIDIA公司推出的一个高性能深度学习推理引擎,它可以将深度学习模型优化为高效的推理引擎,从而在生产环境中实现快速、低延迟的推理。为了方便用户同时使用DeepRec和TensorRT,我们将TensorRT集成进DeepRec,DeepRec会对用户的graph进行分析,将TensorRT能识别的subgraph进行clustering,这里类似XLA的圈图方式。对于每个clustering子图,使用一个TRTEngineOp驱动执行。
环境配置
目前在DeepRec发布的镜像中暂时没有安装TensorRT环境,需要用户手动安装,下一次的镜像发布会带上安装好的环境。 可以自己下载TRT或者使用我们提供的TRT-8.4.2.4
解压tar包并且拷贝到合适的位置,譬如:/usr/local/TensorRT-8.4.2.4,我们需要设置环境变量如下:
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/TensorRT-8.4.2.4/lib
export TENSORRT_INSTALL_PATH=/usr/local/TensorRT-8.4.2.4
生成saved_model
用户需要生成自己的saved model,文档中以"modelzoo/features/embedding_variable/wide_and_deep/train.py"为示例,用户首先需要生成saved model,
python train.py --saved_model_path=./pb
转换模型
对于生成的模型转换成TRT模型,这里需要使用"tensorflow/compiler/tf2tensorrt/tool/tf2trt.py"脚本进行转换。
python tf2trt.py
注意tf2trt.py文件中有一些参数需要自己配置一下,
...
if __name__ == "__main__":
run_params = {
# 模型精度要求,默认是FP32,
# TensorRT 提供了 FP16 量化与 INT8 量化。
'precision_mode':"FP32",
# 是否生成将在运行时构建TRT network和engine的动态TRT ops。
'dynamic_engine': True,
# 对于INT8 量化,TRT需要进行校准(calibration),生成校准文件。
'use_calibration': False,
# input数据的最大batch size
'max_batch_size': 1024,
# 在线转换还是离线转换
'convert_online': False,
# 模型是否使用了embedding variable
'use_ev': True,
}
# 原始的模型saved model所在位置
saved_model_dir = '/model/pb'
# 转换之后的模型保存位置
trt_saved_model_dir = './trtmodel'
ConvertGraph(run_params, saved_model_dir, trt_saved_model_dir)
...
执行模型
使用DeepRec-AI/serving加载新生成的模型。注意需要重新编译serving,并且在此次编译时需要配置环境变量 export TF_NEED_TENSORRT=1
,这样保证在serving so的符号表中是包含了TRT符号的。