Cpp DenseNet Tensorrt导出
Python版本导引
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)