ディープラーニングモデルPyTorchはトレーニングされ、展開のためにONNXとTensorRTに転送されます
一般的に、深層学習モデルはPythonに基づいてトレーニングされます。モデルをデプロイするときは、通常、C ++コードに基づいてデプロイする必要があります。この記事では、例としてLeNetを取り上げ、PyTorchベースのトレーニングからTensorRTのデプロイまで、ディープラーニングモデルをトレーニングするプロセス全体を紹介します。詳細なコードについては、https：//github.com/linghu8812/tensorrt_inference/tree/master/を参照してください。 lenet。
1.モデルトレーニング
python3 train.py手書き数字を認識するモデルはコマンドでトレーニングでき、モデルの認識精度は約99％です。
2.PyTorchモデルをONNXモデルに変換します
python3 export_onnx.pyPyTorchモデルはコマンドを使用してONNXモデルに変換できます。TensorRTフレームワークはPyTorchモデルを直接解析できないため、最初にPyTorchモデルを変換します。次のコードでは、変換のバッチサイズは10です。
import onnx
import torch
# export from pytorch to onnx
net = torch.load('mnist_net.pt').to('cpu')
image = torch.randn(10, 1, 28, 28)
torch.onnx.export(net, image, 'mnist_net.onnx', input_names=['input'], output_names=['output'])
# check onnx model
onnx_model = onnx.load("mnist_net.onnx") # load onnx model
onnx.checker.check_model(onnx_model)
3.lenet_trtをコンパイルします
コードのクローンを作成したら、直接コンパイルします。
mkdir build && cd build
cmake ..
make -j
4.lenet_trtを実行します
コンパイル後、直接実行します./lenet_trt ../config.yaml ../samples/。sampleフォルダには、認識する必要のある写真があります。config.yaml変換されるonnxファイルと変換されたtensorrtファイル名、推論中のバッチサイズ、モデル推論中のチャネルの幅、高さ、数など、モデル操作の構成を定義します。
lenet:
onnx_file: "../mnist_net.onnx"
engine_file: "../mnist_net.trt"
BATCH_SIZE: 10
INPUT_CHANNEL: 1
IMAGE_WIDTH: 28
IMAGE_HEIGHT: 28
操作の結果は次のとおりです。
