[論文レビュー] High Resolution Medical Image Analysis with Spatial Partitioning

Le Hou, Youlong Cheng+7 more|arXiv (Cornell University)|Sep 6, 2019

Medical Image Segmentation Techniques参考文献 18被引用数 29

ひとこと要約

本論文は、Mesh-TensorFlowにおける空間的パーティショニングフレームワークを提示し、512×512×512ボクセルに達する高解像度医用画像上で3D U-Netモデルのエンドツーエンド学習を可能にする。ホールド交換を用いてパーティショニングされたパッチ間の畳み込み接続性を維持する。知る限り、この研究は、ダウンサンプリングやデータのクロッピングなしに、フル解像度のCTスキャンを直接学習する最初の手法である。LiTSベンチマークにおいて、学習時間に5%のオーバーヘッドを伴いながらも、最先端のDiceスコアを達成した。

ABSTRACT

Medical images such as 3D computerized tomography (CT) scans and pathology images, have hundreds of millions or billions of voxels/pixels. It is infeasible to train CNN models directly on such high resolution images, because neural activations of a single image do not fit in the memory of a single GPU/TPU, and naive data and model parallelism approaches do not work. Existing image analysis approaches alleviate this problem by cropping or down-sampling input images, which leads to complicated implementation and sub-optimal performance due to information loss. In this paper, we implement spatial partitioning, which internally distributes the input and output of convolutional layers across GPUs/TPUs. Our implementation is based on the Mesh-TensorFlow framework and the computation distribution is transparent to end users. With this technique, we train a 3D Unet on up to 512 by 512 by 512 resolution data. To the best of our knowledge, this is the first work for handling such high resolution images end-to-end.

研究の動機と目的

• GPU/TPUのメモリ制限により、高解像度医用画像（例：10⁸〜10⁹ボクセルを有する3D CTスキャン）に直接畳み込みニューラルネットワーク（CNN）を学習することが不可能であるという問題に対処すること。
• クロッピングやダウンサンプリング、粗〜細のスケーリング手法といった従来のアプローチに起因する情報損失や複雑な実装の制限を克服すること。
• フル解像度の3D医用画像解析において、効率的かつ透明なモデル並列およびデータ並列を複数のGPU/TPUで実現すること。
• TPUおよびGPU両方の学習をサポートするスケーラブルでオープンソースのフレームワークを、Mesh-TensorFlowに基づいて開発すること。
• 合成ベースのデータオーグメンテーション手法を導入することで、肝腫瘍セグメンテーションの一般化性能を向上させること。

提案手法

• 空間的パーティショニングにより、高解像度3D画像を重複のないパッチに分割し、複数のGPU/TPUに分散配置する。
• 各畳み込み層の前には、隣接するデバイスがパッチの端縁（カーネルサイズの半分）を交換するホールド交換を実行し、空間的文脈を保持する。
• フレームワークはMesh-TensorFlowに基づき、最小限のユーザー介入で自動的なデータ並列およびモデル並列を実現する。
• 合成ベースのデータオーグメンテーション手法により、現実的な強度および形状の変動を持つ合成腫瘍を生成し、モデルのロバスト性を向上させる。
• ダウンサンプリングやパッチベースの推論なしに、512×512×512解像度データ上で3D U-Netモデルのエンドツーエンド学習を可能にする。
• 学習には混合精度（半精度）浮動小数点演算と、Dice損失と交差エントロピー損失の組み合わせを用いたAdafactor最適化手法を採用する。

実験結果

リサーチクエスチョン

• RQ1 ダウンサンプリングやクロッピングなしに、フル解像度の512×512×512 CTスキャン上で3D U-Netモデルをエンドツーエンドで学習可能か？
• RQ2 ホールド交換を伴う空間的パーティショニングにより、高解像度医用画像上で計算オーバーヘッドを最小限に抑えつつ、効率的かつスケーラブルな学習が可能か？
• RQ3 合成ベースのデータオーグメンテーション手法は、3D CTスキャンにおける肝腫瘍セグメンテーションの一般化性能と性能向上に寄与するか？
• RQ4 提案フレームワークを用いた場合、LiTSベンチマークにおけるDiceスコアは、入力解像度が上昇するにつれてどのように変化するか？
• RQ5 分散TPU環境下で、空間的パーティショニングとホールド交換によって実際にどの程度の学習時間オーバーヘッドが生じるか？

主な発見

• 提案フレームワークは、512×512×512解像度のCTスキャン上で3D U-Netモデルの学習に成功し、LiTS検証セットにおいて平均Dice per caseスコアが0.4547 ± 0.0475を達成した。
• 512×512×512解像度で、グローバルDiceスコアは0.7180 ± 0.0446に達し、大規模ボリュームセグメンテーションにおいて優れた性能を示した。
• パーティショニング、リシェーピング、ホールド交換の操作に起因する追加の学習時間オーバーヘッドは約5%にとどまった。
• 提案されたデータオーグメンテーションがなければ、Dice per caseスコアは最低でも10%低下し、モデルの一般化性能においてその重要性が明確になった。
• より高い解像度の入力（例：512³）は、低い解像度（例：64³、128³）よりも一貫して高いDiceスコアを示し、フル解像度学習の利点を裏付けた。
• フレームワークはTPUおよびGPU両方の学習をサポートし、ネットワークアーキテクチャの変更なしにデバイス間での計算を透明に分散可能である。