[論文レビュー] Fast and Accurate 3D Medical Image Segmentation with Data-swapping Method

Haruki Imai, Samuel Matzek+3 more|arXiv (Cornell University)|Dec 19, 2018

Advanced Neural Network Applications被引用数 8

ひとこと要約

本論文は、GPUとCPUメモリ間の途中特徴量を動的に移動させることで、パッチ化せずに192×192×192のフルサイズ3D医療画像上で3D U-Netモデルを訓練可能にするデータスワッピング手法を提案する。この手法により、通信オーバーヘッドが17.1%削減され、全体腫瘍および腫瘍コアのDiceスコアがそれぞれ4.48%および5.32%向上し、訓練時間が164時間から47時間へ3.53倍高速化された。

ABSTRACT

Deep neural network models used for medical image segmentation are large because they are trained with high-resolution three-dimensional (3D) images. Graphics processing units (GPUs) are widely used to accelerate the trainings. However, the memory on a GPU is not large enough to train the models. A popular approach to tackling this problem is patch-based method, which divides a large image into small patches and trains the models with these small patches. However, this method would degrade the segmentation quality if a target object spans multiple patches. In this paper, we propose a novel approach for 3D medical image segmentation that utilizes the data-swapping, which swaps out intermediate data from GPU memory to CPU memory to enlarge the effective GPU memory size, for training high-resolution 3D medical images without patching. We carefully tuned parameters in the data-swapping method to obtain the best training performance for 3D U-Net, a widely used deep neural network model for medical image segmentation. We applied our tuning to train 3D U-Net with full-size images of 192 x 192 x 192 voxels in brain tumor dataset. As a result, communication overhead, which is the most important issue, was reduced by 17.1%. Compared with the patch-based method for patches of 128 x 128 x 128 voxels, our training for full-size images achieved improvement on the mean Dice score by 4.48% and 5.32 % for detecting whole tumor sub-region and tumor core sub-region, respectively. The total training time was reduced from 164 hours to 47 hours, resulting in 3.53 times of acceleration.

研究の動機と目的

• 高解像度の医療画像上で深層3D U-Netモデルを訓練する際のGPUメモリ制限を解消すること。
• パッチベースの訓練手法が原因で生じるオブジェクトのパッチ間での断片化によって引き起こされる品質低下を排除すること。
• GPUとCPUメモリ間の途中活性化を効率的に管理するデータスワッピング技術を開発し、より大きな有効GPUメモリを模倣すること。
• 空間的連続性を損なわず、より高速かつ高精度な3D医療画像セグメンテーションを実現すること。

提案手法

• 訓練中にGPUメモリから中間特徴マップをCPUメモリにオフロードすることで、GPUメモリの圧力を軽減する。
• メモリ使用状況と計算依存関係に基づいて、GPUとCPU間のデータ転送を動的にスケジューリングする。
• 3D U-Netに特化して最適化され、通信オーバーヘッドを最小限に抑えるために、データ転送の頻度とサイズを調整する。
• パッチ化なしに、フル192×192×192ボリューム画像上でエンドツーエンドの訓練が可能となり、空間的文脈が保持される。
• 頻繁にアクセスされる特徴を優先し、重複する転送を最小限に抑えるカスタムメモリ管理戦略を採用する。
• 標準的な指標（Diceスコア、訓練時間など）を用いて、脳腫瘍セグメンテーションデータセット上で評価した。

実験結果

リサーチクエスチョン

• RQ1 データスワッピングは、パッチ化なしにフルサイズ3D医療画像上で3D U-Netを訓練可能であり、セグメンテーション精度を維持または向上させることができるか？
• RQ2 通信オーバーヘッドの観点から、データスワッピングの訓練効率は、標準的なパッチベースの訓練と比較してどの程度優れているか？
• RQ3 パッチ化を回避することで、小規模または断片化した腫瘍領域におけるセグメンテーション性能はどの程度向上するか？
• RQ4 訓練時間の最小化とモデル精度の最大化を実現するためのデータスワッピングパラメータの最適な設定は何か？
• RQ5 パッチベースの代替手法と比較して、データスワッピング手法はより速い収束とより優れた一般化性能を達成できるか？

主な発見

• データスワッピング手法により、ベースライン手法と比較して通信オーバーヘッドが17.1%削減され、訓練効率が顕著に向上した。
• パッチベースの訓練と比較して、フル192×192×192画像での訓練により、全腫瘍サブ領域の平均Diceスコアが4.48%向上した。
• 腫瘍コアサブ領域のDiceスコアは5.32%向上し、重要な領域の局在化がより良好であることが示された。
• 総訓練時間は164時間から47時間に短縮され、訓練速度が3.53倍高速化された。
• 本手法は、精度と速度の両面でパッチベースの訓練を上回り、フルボリュームの文脈を保持することの利点を実証した。
• 結果から、データスワッピングにより、パッチ化に伴うアーティファクトや性能低下を伴わずに高解像度3Dセグメンテーションが可能であることが確認された。