Mesh-TensorFlowによる超高解像度画像解析(2/2)

１．Mesh-TensorFlowによる超高解像度画像解析(2/2)まとめ

・畳み込み演算では各フレームの端部分を超えてフィルターが適用されることがよくある
・ヘィロー交換と呼ばれるデータ通信ステップを実装し畳み込みの前に隣接する部分と通信を行う事で対処
・Mesh-TensorFlowはGPUとTPUの両方で動作し超高解像度画像へのCNNの適用が可能になる

２．Halo Exchangeとは？

以下、ai.googleblog.comより「Ultra-High Resolution Image Analysis with Mesh-TensorFlow」の意訳です。元記事の投稿は2020年2月28日、Le HouさんとYoulong Chengさんによる投稿です。アイキャッチ画像のクレジットはPhoto by Priscilla Du Preez on Unsplash

Halo Exchangeを使用した空間パーティション
画像に対して実行される畳み込み演算では、フレームの端部分を超えてフィルターが適用されることがよくあります。

単一の画像を扱うときにこれに対処する方法はありますが、標準的なアプローチでは、分割された画像が依然としてフレームの端を超えて情報が必要になる可能性がある事を考慮していません。空間的に分割され、各プロセッサに割り当てられた画像の畳み込み演算で正確な結果を得るためには、各分割画像の隣接関係を考慮する必要があります。

潜在的な解決策の1つは、3D画像を分割して出来た各空間パーティションに重複領域を含めることです。ただし、畳み込み層が多数存在し、それぞれが重複する可能性が非常に高いため、重複領域は比較的大きくなります。実際、ほとんどの場合、重複が必要な領域は画像全体とほぼ等しくなります。さらに、重複する領域は全て、あらかじめ最初のレイヤーに含めなければなりません。これは、解決しようとしているメモリの制約にひっかかってしまう可能性があります。

私たちの解決策はまったく異なります。ヘィロー交換(Halo Exchange)と呼ばれるデータ通信ステップを実装しました。すべての畳み込み操作の前に、各空間パーティションはその隣接する部分とマージン部分を交換(受信および送信)し、そのマージン部分で画像を効果的に拡張します。

その後、畳み込み演算は各演算装置で個々に実施されます。これにより、画像全体の畳み込みの結果が、空間分割の有無にかかわらず同一になります。

Halo交換により、パーティション間の畳み込みが分轄された画像を正しく処理することが保証されます。

概念実証-肝腫瘍CTスキャンのセグメンテーション
その後、このフレームワークを肝臓腫瘍の3D CTスキャンタスク(LiTSベンチマーク)に適用しました。評価指標として、ソレンセン-ダイス係数を使用します。これは、「分割された腫瘍領域と真の腫瘍領域の重複がないことを示すスコア0」と「完全な一致を示す1」の間で0.0～1.0の値をとります。以下に示す結果は、より高いデータ解像度がより良い結果をもたらすことを示しています。512の3乗のフル解像度(x、y、z方向のそれぞれで512ピクセル)を使用すると、スコアの伸びが小さくなる傾向がありますが、この手法により超高解像度の画像分析が可能になります。

データの解像度が高いほど、精度が向上します。

結論
既存のデータとモデルの並列化技術により、数十億のパラメーターを使用したニューラルネットワークのトレーニングが可能になりましたが、10の8乗画素以上の入力画像を処理できません。本研究では、これらの超高解像度画像でのCNNの適用性を調査し、有望な結果を示しました。Mesh-TensorFlowベースの実装は、GPUとTPUの両方で動作します。公開れたコードにより、以前は不可能だったいくつかのタスクの可能なソリューションを提供したいと考えています。

謝辞
初期の3D U-netモデルとトレーニングデータを提供してくれたメイヨークリニックの協力者であるPanagiotis Korfiatis博士とDaniel Blezek博士に感謝します。メイヨークリニックとのPOC作業に感謝します。論文のすべての共著者、特にNoam Shazeerの協力に感謝します。