Blog
ブログ

2024年11月20日

YOLOモデルをTensorRTにExportする方法

新人社員のLです。今回はYOLOモデルをTensorRTにExportする方法を紹介します。

まずはAnacondaをダウンロードしてみましょう:

https://www.anaconda.com/download

このサイトに入ってから「skip registration」をクリックしてダウンロードページに入ります。

 

 

するとインストーラーの選択画面が出ますので、
自分のシステムに合わせてInstallerをダウンロードしてください。

インストールする際、Windowsの方はインストール手順に従えば大丈夫です。Linuxの方は、コマンド「bash installer.sh」を実行してください。installerにさきほどダウンロードしたInstallerファイルを置き換えてください。

これでAnacondaはインストールできます。そしてAnaconda3フォルダとその中のScriptsフォルダのパスを環境変数に入れることは忘れないようにしてください。

後、コマンド「conda init」と[conda activate]を実行すれば、anaconda環境が起動されます。(conda init bashなら.bashrcに、conda init zshなら.zshrcに初期化されます。)

YOLOをインストールする前に、一つ新しいAnaconda環境を作ろう。

このコマンド「conda create –name env_name python=version」で新しいAnaconda環境を作れますが、env_nameに環境の名前を置き換えて、そしてversionに指定したいpythonバージョンを置き換えてください。

新しいAnaconda環境が出来た後、「conda activate env_name」を実行すれば、新しい環境に切り替えられます。

次は、Pytorchをインストールする手順です。”https://pytorch.org/

上記のサイトに入って自分のシステムとGPUに適したpytorchをインストールしてください。

(2024/11/19訂正:悲報/朗報?→https://news.yahoo.co.jp/articles/a69637c370074c54ce171d3ba4a8e2f052ab257b、PyTorchが公式の「Anaconda」チャンネルを廃止する予定なので、今後のpytorchインストールはPipを使いましょう。)

 

 

自分のGPUに適したCUDAバージョンは何かというと、”https://developer.nvidia.com/cuda-gpus“、まずはこのサイトで自分のGPUを探し出してください(AMD-YESの方はごめんなさい)。

Compute Capabilityがわかったら、またCUDAの(英語バージョンの!)wikiに入ってください。”https://en.wikipedia.org/wiki/CUDA“。そこには「GPUs supported」というサブタイトルがあって、その下のスプレッドシートを参照してください。そして自分のGPUのCompute Capabilityが緑の区間に当てはまっているものの中で、一番新しいCUDA SDK Versionを選択すれば問題ありません。

例えば、GeForce RTX 4090の場合、Compute Capabilityは8.9だから、CUDA 12.4のpytorchを何も考えずにインストールすればOK。

 

 

Pytorchを さっき作ったAnacondaの仮想環境 にインストールした後、いよいよYOLOをインストールするステップです。

https://github.com/ultralytics/ultralytics“、公式のgithubに参照すると、コマンド「pip install ultralytics」を実行すれば、YOLOのモジュール自体はインストールできます。

今回の記事はトレーニングの記事ではないので、只軽くトレーニングのやり方を説明します。

  1. データセットとデータセットのyamlは必須。”https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/#13-prepare-dataset-for-yolov5“、この記事の”Option 2: Create a Manual Dataset”のところを見てください。
  2. トレーニングのコードサンプルは下記のように。

 

 

トレーニングはここまで。

トレーニングして.ptというモデルの重みファイルを取得すると、それを使ってTensorRTモデルをExportします。詳細はこの記事を参照してください、”https://docs.ultralytics.com/modes/export/“。

簡単に言うと、上記のAnaconda環境の下に下記のコマンドを実行します。

「yolo export model=path/to/model.pt format=engine simplify=True」、modelは.ptファイルのパス、formatは出力のタイプ、simplifyは”Simplifies the model graph”—つまり不要な重みをゼロにすることでスピードアップさせる操作です。

すると、一つの.onnxファイルと.engineファイルが出力されます。.engineファイルを使って、コードに「model =

path/to/model.engine, task=”detect”)」を書けば、普通のpytorchモデルのように使えます。YOLOのpytorchモデルはさらにややこしくなることもあるが、公式のYOLOモデルの使い方を参照すれば無難です。
以上、安定にYOLOモデルをTensorRTにExportする方法でした。なぜAnacondaのインストールまで書くのかも「安定再現」のためです。ネットにいろいろな奇妙なテクニックを記載している記事もあると思いますが、どうやっても不具合に困る場合はこの記事が役に立てば、これを書いた甲斐も多少ありましたね。

このページの先頭へ