[論文レビュー] LAMP: Large Deep Nets with Automated Model Parallelism for Image Segmentation

Wentao Zhu, Can Zhao+4 more|arXiv (Cornell University)|Jun 22, 2020

Advanced Neural Network Applications参考文献 29被引用数 6

ひとこと要約

LAMPは、3D U-NetおよびSEU-Netアーキテクチャを大規模な全画像入力でトレーニング可能な自動モデル並列化を導入し、セグメンテーション精度と推論速度の大幅な向上を実現した。GPipeに基づく並列U-Net設計を活用することで、モデルアーキテクチャの変更を要せず、医療画像セグメンテーションタスクで最先端の性能を達成した。

ABSTRACT

Deep Learning (DL) models are becoming larger, because the increase in model size might offer significant accuracy gain. To enable the training of large deep networks, data parallelism and model parallelism are two well-known approaches for parallel training. However, data parallelism does not help reduce memory footprint per device. In this work, we introduce Large deep 3D ConvNets with Automated Model Parallelism (LAMP) and investigate the impact of both input's and deep 3D ConvNets' size on segmentation accuracy. Through automated model parallelism, it is feasible to train large deep 3D ConvNets with a large input patch, even the whole image. Extensive experiments demonstrate that, facilitated by the automated model parallelism, the segmentation accuracy can be improved through increasing model size and input context size, and large input yields significant inference speedup compared with sliding window of small patches in the inference. Code is available\footnote{https://monai.io/research/lamp-automated-model-parallelism}.

研究の動機と目的

• 3D医療画像セグメンテーションの精度に及ぼすモデルサイズおよび入力コンテキストサイズの増大の影響を調査すること。
• 医療画像分野における大規模3D ConvNetおよび全画像入力のトレーニングを妨げるGPUメモリ制限を解決すること。
• 医療画像解析で広く用いられる標準的なU-Netアーキテクチャと互換性を持つ自動モデル並列化フレームワークを開発すること。
• 大きな受容 field を有する単一パスの全画像推論にスライディングウィンドウ推論を置き換えることで、推論時間を短縮すること。
• カリキュラムトレーニングおよび適応的重み付き損失の有効性を検証し、医療画像セグメンテーションにおけるクラス不均衡の問題を軽減すること。

提案手法

• GPipeのパイプライン並列化に基づく並列U-Netアーキテクチャを設計し、大規模3D ConvNetレイヤーを複数のGPUに分散配置する。
• 手動でのレイヤー分割やアーキテクチャ変更を必要としない、自動モデル並列化を実装する。
• 小規模、中規模、大規模な入力パッチを段階的にトレーニングすることで、大規模モデルの学習を安定化させるカリキュラムトレーニング戦略を採用する。
• Dice損失とフォーカル損失を組み合わせた適応的重み付き損失を適用し、小規模な臓器や病変を有する医療画像のクラス不均衡を是正する。
• 学習率スケジューリングとバッチ正則化を併用したRMSProp最適化手法を用い、大規模モデルのトレーニングを安定化させる。
• 最大512×512×704の大きな入力パッチを用いることで、スライディングウィンドウ推論の必要性を排除し、全画像推論を可能にする。

実験結果

リサーチクエスチョン

• RQ1 3D ConvNetのサイズを拡大すると、医療画像データセットにおけるセグメンテーション精度にどのように影響するか？
• RQ2 スライディングウィンドウ推論と比較して、全画像または大規模パッチを用いた大きな入力コンテキストが、推論速度および精度に与える影響は何か？
• RQ3 自動モデル並列化により、アーキテクチャの変更なしに、標準的な医療画像データセット上で大規模3D U-NetおよびSEU-Netモデルをトレーニング可能か？
• RQ4 入力サイズを段階的に増大させたカリキュラムトレーニングは、大規模モデルの収束性および性能にどのように寄与するか？
• RQ5 大規模モデルおよび大規模入力と組み合わせた適応的重み付き損失は、小構造セグメンテーションにおけるクラス不均衡をどの程度軽減できるか？

主な発見

• 全画像入力のSEU-Net-128モデルはAnatomyNetを上回る優れた性能を示し、大規模モデルおよび大規模コンテキストが精度向上に寄与することを実証した。
• 全画像入力の採用により、1枚あたりの推論時間が1.52秒（U-Net-32で256³）に短縮され、64³パッチを用いたスライディングウィンドウ推論と比較して5.8倍の高速化が達成された。
• 大規模モデル（例：U-Net-128）と大規模入力により、肝臓セグメンテーションでDiceスコア90.99%、腫瘍セグメンテーションで56.48%を達成し、小規模な構成を大幅に上回った。
• 最大入力（256³）では、U-Net-128の推論時間が4.39秒にまで短縮され、前方伝搬回数を減らすことで大規模入力が高速推論を可能にした。
• カリキュラムトレーニング戦略により、段階的に大きなパッチでトレーニングすることで、大規模モデルの安定した学習が可能となり、高い精度が達成された。
• 自動モデル並列化フレームワークは、1ブロックあたり最大128フィルタを有するモデルおよび最大512×512×704の入力サイズまで対応し、GPUメモリ制限を克服した。