onnx转tensorflow以便于发挥显卡的混合精度

triton目前只有tensorflow可以打开混合精度充分使用T4和V100显卡的混合精度计算单元。

#import os
#os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import onnx
from onnx_tf.backend import prepare
onnx_model = onnx.load("onnx/model.onnx")  # load onnx model
#strict需要为False，否则转换的TF模型运行报错
tf_rep = prepare(onnx_model, strict=False)  # prepare tf representation
#tf_folder里的就是TF模型，可直接拷贝到triton里运行
tf_rep.export_graph("tf_folder") # export the model

性能测试

onnx选择sbert的动态轴onnx

TF模型在triton里打开混合精度

测试时，和pytorch一样，triton的perf_analyzer需要使用--input-data zero，否则测试报错

Tensorflow混合精度下的测试

sequence 7

batch 1: 38

batch 50: 1620

batch 100: 2600

sequence 128

batch 1: 38

batch 50: 840

batch 100: 1020

sequence: 512

batch 1: 35

batch 50: 230

batch 100: 240

onnx的测试

当sequence为512时

batch 1: 82

batch 50： 100

测试比较

在单个处理的速度上，onnx胜过使用了混合精度的Tensorflow模型；

在批量处理上，如批量50条，则应用了混合精度的Tensorflow模型比Onnx模型性能高2倍多。