序文
この記事では、トレーニング済みの pytorch モデルをデプロイ用の tensorrt モデルに変換するいくつかの方法を要約します。変換原理のプロセスは大まかに次のとおりです。
- ネットワーク定義と関連する重みをエクスポートします。
- ネットワーク定義と関連する重みを分析します。
- グラフィックス カード オペレーターに基づいて最適な実行計画を構築します。
- 実行計画をシリアル化して保存します。
- 実行計画を逆シリアル化します。
- 推論する
3 番目の点に注目してください。tensorrt によって変換されたモデルが実際にハードウェアにバインドされていることがわかります。つまり、デプロイメント プロセス中に、グラフィック カードおよびグラフィック カード関連のドライバー ソフトウェア (cuda、cudnn) が変更された場合、その場合、モデルを再変換する必要があります。
1. トレテクス
trtexec は、tensorrt パッケージに付属する変換プログラムです。このプログラムは bin ディレクトリにあります。使いやすく、trt モデルを変換する最も簡単な方法です。使用する前に、cuda と cudnn をシステムにインストールする必要があります。そうしないと正常に動作しません。使用例は以下のとおりです。
まず、pytorch モデルを onnx モデルに変換します。サンプルコードは次のとおりです。
def torch2onnx(model_path,onnx_path):
model = load_model(model_path)
test_arr = torch.randn(1,3,32,448)
input_names = ['input']
output_names = ['output']
tr_onnx.export(
model,
test_arr,
onnx_path,
verbose=False,
opset_version=11,
input_names=input_names,
output_names=output_names,
dynamic_axes={
"input":{
3:"width"}} #动态推理W纬度，若需其他动态纬度可以自行修改，不需要动态推理的话可以注释这行
)
print('->>模型转换成功！')
trtexec変換コマンドは以下のとおりです。
固定サイズのモデル変換:
./trtexec --onnx=repvgg_a1.onnx --saveEngine=repvgg_a1.engine --workspace=1024 --fp16
動的サイズモデル変換:
./trtexec --onnx=repvgg_a1.onnx --saveEngine=repvgg_a1.engine --workspace=1024 --minShapes=input:1x3x32x32 --optShapes=input:1x3x32x320 --maxShapes=input:1x3x32x640 --fp16
パラメータの詳細な説明:
- –onnx onnx パス
- –saveEngine trt シリアル化推論エンジンの保存アドレス
- –workspace ワークスペースのサイズをメガバイト単位で設定します (デフォルト = 16)
- –minShapes 提供された最小形状構成ファイルを使用して動的形状を生成します。
- –optShapes 提供された最適な形状の構成ファイルを使用して動的形状を生成します
- –maxShapes 提供された最大形状構成ファイルを使用して動的形状を生成します。
- –fp16 float16 精度推論をオンにします (このモードは推奨されます。一方で速度を上げることができますが、他方では精度の低下は比較的小さいです)
2.トーチ2trt
torch2trt は、nvidia によって公式に保守されている、使いやすい PyTorch to TensorRT コンバータです。使用は比較的簡単ですが、環境設定は上記の方法よりも複雑です。Torch、torch2trt、tensorrt を事前にインストールする必要があります。 Python 環境に tensorrt をインストールするには たとえば、Tensorrt の .tar パッケージでこれらの whl パッケージを見つけ、pip を使用して直接インストールします。
#1、安装tensorrt
cd ~/TensorRT-8.2.4.2/python
pip install tensorrt-8.2.4.2-cp37-none-linux_x86_64.whl
#2、安装Python UFF wheel文件。只有当你将TensorRT与TensorFlow一起使用时才需要安装这个文件 用处：pb转tensorRT
cd ~/TensorRT-8.2.4.2/uff
pip install uff-0.6.9-py2.py3-none-any.whl
#3、安装Python graphsurgeon whl文件 用处：可以让TensorRT 自定义网络结构
cd ~/TensorRT-8.2.4.2/graphsurgeon
pip install graphsurgeon-0.4.5-py2.py3-none-any.whl
#注意trt7.0的版本没有这个包(不用装)
#4、安装Python onnx-graphsurgeon whl文件
cd ~/TensorRT-8.2.4.2/onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl
#5、安装pycuda 可以通过它来实现python 下CUDA 的编程
pip install pycuda
#6、验证安装,打印出tensorrt版本，即安装成功
python
import tensorrt
tensorrt.__version__
torch2trt のインストール:
git clone https://github.com/NVIDIA-AI-IOT/torch2trt.git
cd torch2trt
sudo python setup.py install --plugins
モデル変換コードの使用例：
model = load_model(model_path)
model.cuda()
arr = torch.ones(1, 3, 32, 448).cuda()
model_trt = torch2trt(model,
[arr],
fp16_mode=True,
log_level=trt.Logger.INFO,
max_workspace_size=(1 << 32),
max_batch_size=1,
)
torch.save(model_trt.state_dict(), os.path.join(output, "model_trt.pth"))
logger.info("Converted TensorRT model done.")
engine_file = os.path.join(output, "model_trt.engine")
with open(engine_file, "wb") as f:
f.write(model_trt.engine.serialize())
logger.info("Converted TensorRT model engine file is saved for C++ inference.")
保存された pth モデルは TRTModule にロードでき、TRTModule はトーチ モデルと同様に通常どおり推論できます。
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('model_trt.pth'))
エンジンでシリアル化されたファイルは、tensorrt プログラムの読み込み推論に使用できますが、ここでは torch2trt が動的推論をサポートしていないことに注意してください。その他の使用例については、 torch2trtのgithub の説明を参照してください。
三、torch2trtダイナミック
torch2trt Dynamic はtorch2trt の動的推論バージョンであり、モデル変換後の動的推論をサポートします。基本的に使用中の torch2trt と同じです。最初のステップは以下をインストールすることです。
git clone https://github.com/grimoire/torch2trt_dynamic.git
cd torch2trt_dynamic
python setup.py develop
次に、動的なスケール パラメーターを追加する使用例があります。
from torch2trt_dynamic import torch2trt_dynamic
import torch
from torch import nn
from torchvision.models.resnet import resnet50
import os
# create some regular pytorch model...
model = resnet50().cuda().eval()
# create example data
x = torch.ones((1, 3, 224, 224)).cuda()
# convert to TensorRT feeding sample data as input
opt_shape_param = [
[
[1, 3, 128, 128], # min
[1, 3, 256, 256], # opt
[1, 3, 512, 512] # max
]
]
model_trt = torch2trt_dynamic(model, [x], fp16_mode=False, opt_shape_param=opt_shape_param)
torch.save(model_trt.state_dict(), os.path.join(output, "model_trt.pth"))
logger.info("Converted TensorRT model done.")
engine_file = os.path.join(output, "model_trt.engine")
with open(engine_file, "wb") as f:
f.write(model_trt.engine.serialize())
logger.info("Converted TensorRT model engine file is saved for C++ inference.")
4. パーサーはNXモデルを分析します
ツール変換を使用したくない場合は、独自のコードを作成し、tensorrt のパーサー インターフェイスを使用して onnx モデルを解析し、エンジン エンジンを構築することもできます。この方法は比較的単純で、他のライブラリに依存せず、サポートされています。動的推論モデルの変換 Python コードの例は次のとおりです。
# --*-- coding:utf-8 --*--
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import time
import cv2, os
import numpy as np
import math
TRT_LOGGER = trt.Logger()
class HostDeviceMem(object):
def __init__(self, host_mem, device_mem):
"""
host_mem: cpu memory
device_mem: gpu memory
"""
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
def get_engine(max_batch_size=1, onnx_file_path="", engine_file_path="", fp16_mode=False, save_engine=False,input_dynamic=False):
"""
params max_batch_size: 预先指定大小好分配显存
params onnx_file_path: onnx文件路径
params engine_file_path: 待保存的序列化的引擎文件路径
params fp16_mode: 是否采用FP16
params save_engine: 是否保存引擎
returns: ICudaEngine
"""
# 如果已经存在序列化之后的引擎，则直接反序列化得到cudaEngine
if os.path.exists(engine_file_path):
print("Reading engine from file: {}".format(engine_file_path))
with open(engine_file_path, 'rb') as f, \
trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read()) # 反序列化
else: # 由onnx创建cudaEngine
# 使用logger创建一个builder
# builder创建一个计算图 INetworkDefinition
explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
# In TensorRT 7.0, the ONNX parser only supports full-dimensions mode, meaning that your network definition must be created with the explicitBatch flag set. For more information, see Working With Dynamic Shapes.
with trt.Builder(TRT_LOGGER) as builder, \
builder.create_network(explicit_batch) as network, \
trt.OnnxParser(network, TRT_LOGGER) as parser: # 使用onnx的解析器绑定计算图，后续将通过解析填充计算图
# builder.max_workspace_size = 1 << 30 # 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30
builder.max_batch_size = max_batch_size # 执行时最大可以使用的batchsize
if fp16_mode:
config.set_flag(trt.BuilderFlag.FP16)
# builder.fp16_mode = fp16_mode
# 解析onnx文件，填充计算图
if not os.path.exists(onnx_file_path):
quit("ONNX file {} not found!".format(onnx_file_path))
print('loading onnx file from path {} ...'.format(onnx_file_path))
with open(onnx_file_path, 'rb') as model: # 二值化的网络结果和参数
print("Begining onnx file parsing")
parser.parse(model.read()) # 解析onnx文件
# parser.parse_from_file(onnx_file_path) # parser还有一个从文件解析onnx的方法
print("Completed parsing of onnx file")
# 填充计算图完成后，则使用builder从计算图中创建CudaEngine
print("Building an engine from file{}' this may take a while...".format(onnx_file_path))
if input_dynamic: # 动态推理
profile = builder.create_optimization_profile()
profile.set_shape("input",(1,3,32,32),(1,3,32,320),(1,3,32,640))
config.add_optimization_profile(profile)
#################
print(network.get_layer(network.num_layers - 1).get_output(0).shape)
engine = builder.build_engine(network, config)
print("Completed creating Engine")
if save_engine: # 保存engine供以后直接反序列化使用
with open(engine_file_path, 'wb') as f:
f.write(engine.serialize()) # 序列化
return engine
if __name__== "__main__":
# These two modes are depend on hardwares
fp16_mode = True
max_batch_size = 1
onnx_model_path = "./repvgg_a1.onnx"
trt_engine_path = "./repvgg_a1.engine"
# Build an cudaEngine
engine = get_engine(max_batch_size, onnx_model_path, trt_engine_path, fp16_mode,True,True)
C++ バージョン解析 onnx コード例:
//step1:创建logger：日志记录器
class Logger : public ILogger
{
void log(Severity severity, const char* msg) override
{
// suppress info-level messages
if (severity != Severity::kINFO)
std::cout << msg << std::endl;
}
} gLogger;
//step2:创建builder
IBuilder* builder = createInferBuilder(gLogger);
//step3:创建network
const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
//step4:创建parser
nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);
//step5:使用parser解析模型填充network
const char* onnx_filename="./model.onnx"
parser->parseFromFile(onnx_filename, ILogger::Severity::kWARNING);
for (int i = 0; i < parser.getNbErrors(); ++i)
{
std::cout << parser->getError(i)->desc() << std::endl;
}
//step6:标记网络输出
for (auto &s : OUTPUT_BLOB_NAMES)
network->markOutput(*blobNameToTensor->find(s.c_str()));
//step7:创建config并设置最大batchsize和最大工作空间
IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxBatchSize(maxBatchSize);//设置最大batchsize
config->setMaxWorkspaceSize(1 << 30);//2^30 ,这里是1G
//step8:创建engine
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
assert(engine);
//step9:序列化保存engine到planfile
IHostMemory *serializedModel = engine->serialize();
assert(serializedModel != nullptr)
std::ofstream p("xxxxx.engine");
p.write(reinterpret_cast<const char*>(serializedModel->data()), serializedModel->size());
//step10:释放资源
serializedModel->destroy();
engine->destroy();
parser->destroy()
network->destroy();
config->destroy();
builder->destroy();
5. tensorrtx
tensorrtxのモデル構築方法はかなり変わっていて、まず tensorrt 独自の API を使ってネットワークを構築し、重みを与えるというやり方で、変換時にネットワークがしっかり確立していれば、基本的に変換は問題ありませんこれは良いことです。onnx を trt に変換するプロセスで一部の演算子がサポートされていないという問題は解決されますが、プロセスは比較的複雑であり、動的スケール推論はサポートされていません。現在、trt は onnx を非常によくサポートしており、基本的にすべての演算子をサポートしていますonnxモデルは変換できるので、面倒でなければ試してみてください。
tensorrtx のクローンを作成する
git clone https://github.com/wang-xinyu/tensorrtx.git
yolov5.wts ファイルを生成し、ウェイト ファイル yolov5s.pt をダウンロードし、tensorrtx/yolov5/gen_wts.py を Ultralytics/yolov5 にコピーして実行します。
python gen_wts.py
tensorrtx/yolov5 をコンパイルし、yolov5.engine ファイルを生成します。
mkdir build
cd build
cmake ..
make
デフォルトでは、s-model と fp16 推論が生成されます。バッチ 1 のエンジンと同様に、yolov5 の他のモデルでもコード内の関連パラメーターを変更できます。
#define USE_FP16
#define DEVICE 0 // GPU ID
#define NMS_THRESH 0.4
#define CONF_THRESH 0.5
#define BATCH_SIZE 1
#define NET s // s m x l
ファイル yolov5.wts を tensorrtx/yolov5/build ディレクトリにコピーし、次のコマンドを実行して yolov5.engine を生成します。
sudo ./yolov5 -s yolov5s.wts yolov5.engine s
sudo ./yolov5 -d yolov5s.engine ../samples
6.onnx-tensorrt
onnx-tensorrt はonnx の公式変換ウェアハウスであり、Tensorrt バージョンの対応するブランチを多数提供しています。たとえば、ここでは 8.2-EA を使用します。正しいコンパイル方法は次のとおりです。
まず、onnx-tensorrt を起動します。
git clone --recursive -b 8.2-EA https://github.com/onnx/onnx-tensorrt.git
コンパイル：
cd onnx-tensorrt
mkdir build
cd build
# /path/to/TensorRT-8.2.4.2改成自己的TensorRT绝对路径
cmake .. -DTENSORRT_ROOT=/path/to/TensorRT-8.2.4.2
make -j8
make install
コンパイルが完了した後、cuda 環境変数が設定されている場合は、再度設定する必要はありません。設定されていない場合は、次のように設定する必要があります。
ターミナルに「vim ~/.bashrc」と入力します。
#cuda
export PATH=/usr/local/cuda-11.4/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
保存して終了し、source ~/.bashrc を実行して更新して有効にします。
onnx-tensorrt 変換コマンドは次のとおりであり、エンジンにシリアル化されます。
onnx2trt my_model.onnx -o my_engine.trt
判読可能な txt テキストに変換します。
onnx2trt my_model.onnx -t my_model.onnx.txt
Python 側で onnx-tensorrt を使用します。
#安装tensorrt
python3 -m pip install <tensorrt_install_dir>/python/tensorrt-8.x.x.x-cp<python_ver>-none-linux_x86_64.whl
#安装onnx
python3 -m pip install onnx==1.8.0
#安装onnx-tensorrt，在onnx-tensorrt目录下运行
python3 setup.py install
推論のための Python コードの使用例:
import onnx
import onnx_tensorrt.backend as backend
import numpy as np
model = onnx.load("/path/to/model.onnx")
engine = backend.prepare(model, device='CUDA:0')
input_data = np.random.random(size=(32, 3, 224, 224)).astype(np.float32)
output_data = engine.run(input_data)[0]
print(output_data)
print(output_data.shape)