Skip to main content

pth2onnx: Convert PyTorch model to Onnx model.

Project description

Convert PyTorch model to Onnx model

PyTorchで作成した重みファイルをONNX形式の重みファイルに変換するCLIアプリケーションです。

サポートしているAIタスクは以下のとおりです。

  • Image Classification
  • Object Detection

動作確認OS

  • Windows 10 Pro
  • Windows 11 Pro

インストール方法

pip install pth2onnx

pth2onnxの実行方法

YOLOXモードの場合

# githubからソースをダウンロードし、YOLOXフォルダ内に.venvを生成しインストールする
pth2onnx -m yolox -c install -f
# Windows環境の場合下記のエラーが出ることがある
# ERROR: Could not install packages due to an OSError: [WinError 206] ファイル名または拡張子が長すぎます。
# これが出たときは「YOLOX/requirements.txt」ファイルの「onnx-simplifier」をコメントアウトして再実行する

# 学習済みモデルのダウンロード先URLを表示
pth2onnx -m yolox -c zoo -f
# see: https://github.com/Megvii-BaseDetection/YOLOX/#benchmark

# pytorchの重みファイルでデモを実行
pth2onnx -m yolox -c demo -f --yolox_model_name <モデル名> --yolox_weight_file <pytorchモデルファイルのパス> --yolox_output_preview
# モデル名は「yolox_nano」「yolox_tiny」「yolox_s」「yolox_m」「yolox_l」「yolox_x」など
# pytorchモデルファイルのパスはYOLOXフォルダ内のパス。「models/yolox_tiny.pth」など

# pytorchの重みファイルをONNXの重みファイルに変換
pth2onnx -m yolox -c convert -f --yolox_model_name <モデル名> --yolox_weight_file <pytorchモデルファイルのパス> --yolox_onnx_file <ONNXモデルファイルのパス>
# ONNXモデルファイルのパスはYOLOXフォルダ内のパス。「models/yolox_tiny.onnx」など

# ONNXの重みファイルで推論を実行
pth2onnx -m yolox -c inference -f --yolox_onnx_file <ONNXモデルファイルのパス> --yolox_model_img_size <モデルのINPUTサイズ> --yolox_output_preview
# モデルのINPUTサイズは「416」など

その他便利なオプション

コマンドラインオプションが多いので、それを保存して再利用できるようにする

# 通常のコマンドに「-u」と「-s」オプションを追加する
pth2onnx -u <オプションを保存するファイル> -s

# 次から使用するときは「-u」を使用する
pth2onnx -u <オプションを保存するファイル>

コマンドの実行結果を見やすくする。

# 通常のコマンドに「-f」オプションを追加する
pth2onnx -f

# 「-f」オプションを外せば、結果はjson形式で取得できる
pth2onnx 

コマンドラインオプションのヘルプ。

pth2onnx -h

pth2onnxコマンドについて

python -m pth2onnxの省略形です。 実体はscriptsディレクトリ内にあります。

データの保存場所

pathlib.Path(HOME_DIR) / '.pth2onnx'

動作確認したモデル

AI Task base Model
Object Detection YOLOX YOLOX-Nano
Object Detection YOLOX YOLOX-Tiny
Object Detection YOLOX YOLOX-s
Object Detection YOLOX YOLOX-m
Object Detection YOLOX YOLOX-l
Object Detection YOLOX YOLOX-x

開発環境構築

git clone https://github.com/hamacom2004jp/pth2onnx.git
cd pth2onnx
python -m venv .venv
.venv\Scripts\activate
python.exe -m pip install --upgrade pip
pip install -r requirements.txt
deactivate

pyplにアップするための準備

python setup.py sdist
python setup.py bdist_wheel
[distutils]
index-servers =
  pypi
  testpypi

[pypi]
repository: https://upload.pypi.org/legacy/
username: __token__
password: 本番環境のAPIトークン

[testpypi]
repository: https://test.pypi.org/legacy/
username: __token__
password: テスト環境のAPIトークン
  • テスト環境にアップロード .pyplrcを作っていない場合はコマンド実行時にusernameとpasswordを要求される 成功するとURLが返ってくる。
twine upload --repository testpypi dist/*
  • pipコマンドのテスト
pip install -i https://test.pypi.org/simple/ pth2onnx
  • 本番環境にアップロード
twine upload --repository pypi dist/*

Lisence

This project is licensed under the MIT License, see the LICENSE.txt file for details

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

pth2onnx-0.0.4.tar.gz (14.0 kB view hashes)

Uploaded Source

Built Distribution

pth2onnx-0.0.4-py3-none-any.whl (13.7 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page