Hugging Face x ONNXを調べてみた〜学習編〜

お久しぶりです。気づいたら前回の投稿から1ヶ月以上も経っていました。。。もう7月も終わりなんて、信じられないですね😅

今回は、以前から気になっていたHugging FaceにおけるONNX(オニキス)の活用方法について調べてみました。きっかけは、以下の公式ブログでHugging Faceのモデルを結構簡単にONNX形式に変換できるんだなぁと思ったことで、kaggleなどで活用できるように理解を深めたいというモチベーションが湧いてきたからです。

Convert Transformers to ONNX with Hugging Face Optimum

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

実際に調べてみると、ONNX形式に変換して推論を行うだけでなく、ONNX Runtimeを使用して学習できることがわかったため、学習編と推論編に分けてご紹介していきます。結論から言うと、今回の実験設定ではONNX Runtimeを使うことで高速化はできなかったのですが、興味のある方はぜひ最後までご覧ください。（A100を使った実験では割と高速化できているそうです。→https://github.com/huggingface/optimum/tree/main/examples/onnxruntime/training/text-classification）

そもそもONNX, ONNX Runtimeとは？
ONNX Runtimeを用いた学習について
おわりに

そもそもONNX, ONNX Runtimeとは？

ONNXの正式名称は、Open Neural Network Exchangeで、学習したモデルなどを異なるフレームワークで共通して利用できるようにするライブラリです。このように、様々なライブラリがサポートされています。（GitHub: https://github.com/onnx/onnx）

ONNX | Supported Tools

ONNX Runtimeは、ONNX形式での学習や推論を行う際に、ハードウェアの性能を最大限発揮できるように最適化を行ってくれるライブラリです。PyTorchとTensorFlowのモデルを両方使いたいケースなどでは、ONNX形式に統一して実行するとコードをシンプルにしつつ、実行速度も担保してくれるのかもしれません。（GitHub: https://github.com/microsoft/onnxruntime）

ONNX Runtimeを用いた学習について

今回は、Hugging Faceのoptimumというライブラリに実装されているORTTrainerを使ってみました。GitHubのREADMEによると、transformersのTrainerをORTTrainerに置き換えるだけでいいとのことで、すごく簡単に試すことができました。

GitHub - huggingface/optimum: 🏎️ Accelerate training and inference of 🤗 Transformers with easy to use hardware optimization tools

🏎️ Accelerate training and inference of 🤗 Transformers with easy to use hardware optimization tools - GitHub - huggingface/optimum: 🏎️ Accelerate training and i...

ここから、実際に動かすために行ったことを紹介していきます。動作確認は、kaggleのnotebookで行っています。

ライブラリのインストール（やや苦労した）

まずは、READMEに従ってoptimumをインストールします。（執筆時点でのライブラリの最新verは1.3.0ですが、ORTTrainingArgumentsを使いたかったので、git経由でインストールしました。）

python -m pip install git+https://github.com/huggingface/optimum.git#egg=optimum[onnxruntime]

GPU環境で動かしたかったので、onnxruntime-gpuもインストールしておきました。

pip install onnxruntime-gpu

READMEによると、必要なインストール作業はこれで完了なのですが、いざORTTrainerを動かしてみるとtorch-ortがない というエラーが出てしまいました・・・

そこで、optimumのexamplesを見に行くと、READMEの下の方に小さく書かれていました。（別でインストールが必要なようです）

torch-ortのインストールに関してですが、PyPIのページを参考に行いました。kaggle環境は、執筆時点でCUDAのバージョンが11.4だったので、以下のコマンドでインストールに成功しました。

# CUDA Version: 11.4におけるonnxruntime-trainingのインストール
pip install onnxruntime-training -f https://download.onnxruntime.ai/onnxruntime_stable_cu114.html
# torch-ortのインストール
pip install torch-ort
# post-installation step
python -m torch_ort.configure

これで、実行できると思ったら、次は ImportError: cannot import name 'PropagateCastOpsStrategy' from 'onnxruntime.capi._pybind_state' (/opt/conda/lib/python3.7/site-packages/onnxruntime/capi/_pybind_state.py) というエラーが出てしまいました・・・

このエラーは、Google先生に聞いたところ、このPRにたどり着き、以下の１行を実行することで解消できました。

sudo apt install libgomp1

実験設定について

今回は、以前参加したNBMEコンペで使用していたTrainerを用いたコードで実験を行いました。

実行環境：kaggle notebook（GPUはP100）
モデル：DeBERTa-V3-large
データ数：5fold（最初のfold0のみ実行）
Sequence bucketing（バッチごとにpadding）を使用
その他のパラメータ：
- epoch: 3
- train_batch_size: 6
- valid_batch_size: 32
- lr: 3e-5
- eval_steps: 500（500ステップごとに評価）
- fp16: True

transformersのバージョンに関しては、執筆時点で最新の4.21.0を使用しました。というのも、DeBERTa-V2アーキテクチャでONNXがサポートされたのが4.21.0からで、これはかなりの朗報なのではないでしょうか！これで、主要なモデルは全てONNX対応した気がします。（リリースノート）

実験結果について

必要なライブラリをインストールし、TrainerをORTTrainer、TrainingArgumentsをORTTrainingArgumentsに置き換えたものとオリジナルのものをそれぞれ実行し、学習部分の時間を比較してみました。

ほとんど変わらないですね・・・（何なら遅くなってる😭）冒頭でも述べましたが、公式のexamplesでは割と高速化しているので、GPUのアーキテクチャに依存する部分が大きいのかもしれないですね。（ちゃんとした原因を特定するには、もう少し調査が必要そうです）

おわりに

今回は、残念ながら学習の高速化は実現できなかったのですが、気軽にONNX Runtimeを用いた学習を試すことができて良かったです。Hugging Face Trainerをお使いの方は、少しの変更でトライできるものなので、ぜひ手元のコードで試してみていただけると幸いです。（もし、高速化に成功したよ！という方がいらっしゃれば、コメント等いただけるとすごく嬉しいです。）

少し話が変わりますが、NBMEコンペからTrainerを愛用するようになったので、基本的な機能紹介やkaggleで使う際のTips的な部分を、今後はご紹介できればと考えています。最後まで、ご覧いただきありがとうございました。