open-mmlab/mmengine

OpenMMLab Foundational Library for Training Deep Learning Models

解決する課題

MMEngine は、PyTorch におけるディープラーニングモデルのトレーニング用に、標準化された基盤フレームワークを提供します。トレーニングループ、検証、モニタリングのための繰り返しのボイラープレートコードを書く必要がなくなるため、研究者や開発者はインフラストラクチャではなくアルゴリズム設計に集中できます。

仕組み

トレーニングエンジンとして機能し、以下の主要コンポーネント間の連携をオーケストレーションします:

  • Runner: トレーニングと検証プロセスを管理する中心的コーディネーター。
  • Model: BaseModel クラスで、モデルのトレーニング（ロス計算）と予測方法を定義します。
  • Dataset & DataLoader: データをモデルに供給する標準化された手段。
  • Metrics: BaseMetric を用いた、トレーニング中およびトレーニング後のモデル性能評価システム。
  • OptimWrapper: オプティマイザをラップし、逆伝播と勾配更新を処理します。

対象者

PyTorch モデルを構築するディープラーニング研究者やエンジニア向けに設計されています。特に OpenMMLab エコシステム内のユーザー、あるいは自分自身の非 OpenMMLab プロジェクト向けに汎用的で設定可能なトレーニングエンジンを求める方に適しています。

ハイライト

  • 大規模トレーニングサポート: DeepSpeed、ColossalAI、FSDP などの主流フレームワークと統合。
  • トレーニング戦略: 混合精度トレーニング、勾配蓄積、勾配チェックポイントなど、メモリと速度を最適化する組み込みサポート。
  • 柔軟な設定: 純粋な Python スタイルとプレーンテキスト（JSON/YAML）設定ファイルの両方に対応。
  • 幅広いモニタリング統合: TensorBoard、WandB、MLflow、ClearML、Neptune、DVCLive、Aim など多数の可視化プラットフォームと接続可能。