Cpp DenseNet Tensorrt导出

Cpp DenseNet Tensorrt导出

Python版本导引

• Python DenseNet OpenVino导出
• Python DenseNet OpenVino推理
• Python DenseNet OpenVino打包

Cpp OpenVINO版本导引

• Cpp DenseNet OpenVino导出
• Cpp DenseNet OpenVino测试
• Cpp DenseNet OpenVino环境配置
• Cpp DenseNet OpenVino CMake工程
• Cpp DenseNet OpenVino辅助函数
• Cpp DenseNet OpenVino部署主体

模型导出

Tensorrt不能直接识别Torch模型，所以需要先导出为Onnx在进行转换

Torch -> Onnx

# encoding=utf-8
# https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html

import torch
import torch.nn as nn
from torchvision import models

# load model
model = models.densenet121(pretrained=True)
model.classifier = nn.Linear(model.classifier.in_features, 5)
model.load_state_dict(torch.load("state.pth"))
model.eval()

# convert by onnx
torch.onnx.export(
    model,
    torch.randn(1, 3, 224, 224),
    "export_dense121_cpu.onnx",
    verbose=True,
    opset_version=16,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output']
)

# test onnx
import onnxruntime
session = onnxruntime.InferenceSession("export_dense121_cpu.onnx")
inputs = {"input":  torch.randn(1, 3, 224, 224).numpy()}
out = session.run(None, inputs)
print(out)

onnx -> engine

Tensorrt和OpenVINO不一样，他不会导出一个中间描述IR模型，它会直接编译当前模型到显卡上，所以这个框架即是一个中介人，也是一个优化器

注意这种模式虽然提高了运行效率，但是移植性会大大降低，通常导出的模型只能在相同型号的显卡上运行，举个例子：在RTX 3090上完成开发流程，但是部署到RTX 2060上会直接不识别

# encoding=utf-8
# https://github.com/NVIDIA/TensorRT/blob/main/samples/python/introductory_parser_samples/onnx_resnet50.py
# https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/gettingStarted.html
# https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Core/BuilderConfig.html?highlight=memorypooltype

# 10.x api is not the same, rebuild
import tensorrt as trt

print("trt version: ", trt.__version__)
TRT_LOGGER = trt.Logger(trt.Logger.INFO)

builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(0)
config = builder.create_builder_config()
parser = trt.OnnxParser(network, TRT_LOGGER)

# build flag
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 * 1 << 30)  # 1GB
#config.set_flag(trt.BuilderFlag.FP16)

with open("export_dense121_cpu.onnx", 'rb') as model:
    if not parser.parse(model.read()):
        print("ERROR: Failed to parse the ONNX file.")
        for error in range(parser.num_errors):
            print(parser.get_error(error))
        exit(-1)


engine_bytes = builder.build_serialized_network(network, config)
with open("export_dense121_gpu.engine", 'wb') as engine:
    engine.write(engine_bytes)