TensorRT教程9：使用 Python parser API解析模型（重点）

使用 Python parser API解析模型（重点）

1、解析模型创建engine的10个基本步骤

step1：创建logger

step2：创建builder

step3：创建network

step4：创建parser

step5：使用parser解析模型填充network对象

step6：标记网络输出

step7：创建config并设置最大batchsize和最大工作空间

step8：创建engine

step9：序列化保存engine

step10：释放资源

2、parser解析caffe模型

#导入模块
import tensorrt as trt
#定义数据类型
datatype = trt.float32
#定义caffe模型路径
deploy_file = 'data/mnist/mnist.prototxt' #网络文件
model_file = 'data/mnist/mnist.caffemodel'#模型文件

#step1：创建logger：日志记录器, 此处我们抑制了信息 消息，并仅报告警告和错误
logger = trt.Logger(trt.Logger.WARNING)

#step2：创建builder、network、parser
with trt.Builder(logger) as builder, builder.create_network() as network, trt.CaffeParser() as parser:
    model_tensors = parser.parse(deploy=deploy_file, model=model_file, network=network, dtype=datatype)

#step3：创建config并设置最大batchsize和最大工作空间
with builder.create_builder_config() as config
    config.max_batch_size= 32
    config.max_workspace_size = 10 << 20
    
#step4：创建engine
engine = builder.build_cuda_engine(network，config)

#step5:序列化保存engine到planfile
with open(“sample.engine”, “wb”) as f:
    f.write(engine.serialize())

3、parser解析onnx模型

在 TensorRT 7.0 中，ONNX 解析器仅支持全维模式，这意味着您的网络定义必须使用 explicitBatch 标志集创建

#导入模块
import tensorrt as trt
#定义数据类型
datatype = trt.float32
#定义onnx模型路径
model_file = 'data/mnist/mnist.onnx'#模型文件

#step1：创建logger：日志记录器, 此处我们抑制了信息 消息，并仅报告警告和错误
logger = trt.Logger(trt.Logger.WARNING)

#step2：创建builder、network、parser
#创建的是动态shape网络
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(logger) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network, logger) as parser:
with open(model_path, 'rb') as model:
    if not parser.parse(model.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))

#step3：创建config并设置最大batchsize和最大工作空间
with builder.create_builder_config() as config
    config.max_batch_size= 32
    config.max_workspace_size = 10 << 20
    
#step4：创建engine
engine = builder.build_cuda_engine(network，config)

#step5:序列化保存engine到planfile
with open(“sample.engine”, “wb”) as f:
    f.write(engine.serialize())

4、parser解析UFF模型（tensorflow模型）

首先使用 UFF 转换器转换冻结 tensorflow的张量流模型到 UFF 文件

convert-to-uff frozen_inference_graph.pb

#导入模块
import tensorrt as trt
#定义数据类型
datatype = trt.float32
#定义uff模型路径
model_file = '/data/mnist/mnist.uff' 

#step1：创建logger：日志记录器, 此处我们抑制了信息 消息，并仅报告警告和错误
logger = trt.Logger(trt.Logger.WARNING)

#step2：创建builder、network、parser
with trt.Builder(logger) as builder, builder.create_network() as network, trt.UffParser() as parser:
    #注册网络输入输出
    parser.register_input("Placeholder", (1, 28, 28))
    parser.register_output("fc2/Relu")
    parser.parse(model_file, network)

#step3：创建config并设置最大batchsize和最大工作空间
with builder.create_builder_config() as config
    config.max_batch_size= 32
    config.max_workspace_size = 10 << 20
    
#step4：创建engine
engine = builder.build_cuda_engine(network，config)

#step5:序列化保存engine到planfile
with open(“sample.engine”, “wb”) as f:
    f.write(engine.serialize())