A Partial Replication of MaskFormer in TensorFlow on TPUs for the TensorFlow Model Garden

要約

このペーパーでは、もともと PyTorch フレームワークを使用して開発されたユニバーサル画像セグメンテーション モデルである MaskFormer モデルを TensorFlow エコシステム内で複製し、特に Tensor Processing Unit (TPU) での実行用に最適化するタスクに取り組みます。
私たちの実装は、TensorFlow Model Garden (TFMG) 内で利用可能なモジュール構造を活用しており、MaskFormer モデルの仕様を満たすように調整および適合された、データ ローダー、トレーニング オーケストレーター、さまざまなアーキテクチャ コンポーネントなどの要素を網羅しています。
私たちは、レプリケーション中に遭遇する主要な課題、非収束の問題、トレーニングの遅さ、損失関数の適応、TPU 固有の機能の統合に対処します。
再現された実装を検証し、COCO データセットに関する定性的な結果を示します。
私たちの実装はエンドツーエンドの再現性に関する目標の一部を満たしていますが、TensorFlow で MaskFormer の PyTorch バージョンを複製する際に課題に直面しました。
この複製プロセスは単純ではなく、多大なエンジニアリング作業が必要です。
具体的には、徹底的な検証とハイパーパラメーターの調整に加えて、TFMG 内のさまざまなコンポーネントのカスタマイズが必要です。
レプリケーションは、https://github.com/PurdueDualityLab/tf-maskformer/tree/main/official/projects/maskformer で入手できます。

要約(オリジナル)

This paper undertakes the task of replicating the MaskFormer model a universal image segmentation model originally developed using the PyTorch framework, within the TensorFlow ecosystem, specifically optimized for execution on Tensor Processing Units (TPUs). Our implementation exploits the modular constructs available within the TensorFlow Model Garden (TFMG), encompassing elements such as the data loader, training orchestrator, and various architectural components, tailored and adapted to meet the specifications of the MaskFormer model. We address key challenges encountered during the replication, non-convergence issues, slow training, adaptation of loss functions, and the integration of TPU-specific functionalities. We verify our reproduced implementation and present qualitative results on the COCO dataset. Although our implementation meets some of the objectives for end-to-end reproducibility, we encountered challenges in replicating the PyTorch version of MaskFormer in TensorFlow. This replication process is not straightforward and requires substantial engineering efforts. Specifically, it necessitates the customization of various components within the TFMG, alongside thorough verification and hyper-parameter tuning. The replication is available at: https://github.com/PurdueDualityLab/tf-maskformer/tree/main/official/projects/maskformer

arxiv情報

提供元, 利用サービス

arxiv.jp, Google