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

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

・学習手法の進歩によりCNNモデルは数分で数百万の画像のデータセットを使って学習可能になった
・しかしCTスキャンなどの三次元画像は二次元画像に比べて画素数が非常に多いので依然として取り扱い困難
・Mesh-TensorFlowはこの超高解像度画像問題を効率的にデータを分割して処理する事に解決する

２．Mesh-TensorFlowとは？

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

ディープニューラルネットワークモデルは、多くの最先端の画像分析技術と自然言語処理アルゴリズムの屋台骨を形成しています。データやモデルを並列処理する手法など、最近の大規模深層学習手法の進歩により、巨大な畳み込みニューラルネットワーク(CNN)モデルは数分で数百万の画像のデータセットを使って学習可能です。

ただし、画素数が最大で10の8乗にまで達する可能性がある3Dコンピューター断層(CT:Computed Tomography)撮影画像などの超高解像度画像にCNNモデルを適用することは依然として困難です。

個々のGPUまたはTPUには通常12～32GBのメモリしかありませんが、既存の手法では、プロセッサは少なくとも32GBの部分的な中間データを保持する必要があります。

典型的な解決策は、画像を細かく分割して別々に処理することです。しかしこれにより、プログラムが複雑になり、情報が損なわれ、パフォーマンスは最適ではなくなります。

論文、「High Resolution Medical Image Analysis with Spatial Partitioning」では、Mayo Clinicとのコラボレーションにより、Mesh-TensorFlowフレームワークを使用して大量データとモデルの並列性の限界を押し広げ、実現を優先して実用的な入力解像度を損なう事なくこの手法を超高解像度画像分析に使用する方法を示します。

隣接する分割データ間の関係を保持するために、ヘィロー交換アルゴリズム(訳注：halo exchange algorithm、直訳すると後光。巨大データをAとBに分割して処理する際は、AとBの境目の部分を互いに参照する必要が出てくる事があります。重複する部分は外周部分なので後光の部分にも見えるのでhaloって事なんだと思われます)を実装して、空間パーティション全体の畳み込み演算を処理しました。

その結果、モデル並列処理を256にして超高解像度画像(各次元が512画素を持つ3D画像)で3D U-Netをトレーニングできました。更に、GPUおよびTPU向けのMesh-TensorFlowベースのフレームワークを、より広範な研究コミュニティで使用するためにオープンソース化しました。

Mesh-TensorFlowを使用したデータとモデルの並列処理
実装は、Mesh-TensorFlowフレームワークに基づいており、簡単で効率的なデータとモデルの並列性を実現します。これにより、ユーザー定義の画像レイアウトに従って、デバイスのメッシュ間でテンソルを分割できます。

例えば、ユーザーは計算デバイスのメッシュを16行 x 16列で定義し、256個のプロセッサーに設定できます。(プロセッサーあたり2コア)次に、レイアウトを定義して、画像の空間次元xをプロセッサ行にマッピングし、空間次元yをプロセッサ列にマッピングし、バッチ次元(つまり、同時に処理される画像数)をコアの数にマッピングします。

トレーニングバッチの分割と配布は、ユーザーが実装の詳細を気にすることなく、テンソルレベルでMesh-TensorFlowによって実装されます。次の図は、概念と簡略化された例を示しています。

超高解像度画像の空間分割、上記例は3D CTスキャン画像です。