[論文レビュー] OSS-Net: Memory Efficient High Resolution Semantic Segmentation of 3D Medical Data

Christoph Reich, Tim Prangemeier+2 more|arXiv (Cornell University)|Oct 20, 2021

Radiomics and Machine Learning in Medical Imaging被引用数 7

ひとこと要約

OSS-Net は、局所観測と事前分布に基づく推論戦略を組み込んだ、メモリ効率が良く高解像度な3次元画像セマンティックセグメンテーションフレームワークであり、オカピシティネットワーク（O-Net）を拡張したものである。3D U-Netベースラインと同等のセグメンテーション精度を達成しながら、顕著に少ないメモリを消費し、より高速な推論を可能にすることで、コンsumer GPU 上でも高解像度3次元セグメンテーションが現実可能となる。

ABSTRACT

Convolutional neural networks (CNNs) are the current state-of-the-art meta-algorithm for volumetric segmentation of medical data, for example, to localize COVID-19 infected tissue on computer tomography scans or the detection of tumour volumes in magnetic resonance imaging. A key limitation of 3D CNNs on voxelised data is that the memory consumption grows cubically with the training data resolution. Occupancy networks (O-Nets) are an alternative for which the data is represented continuously in a function space and 3D shapes are learned as a continuous decision boundary. While O-Nets are significantly more memory efficient than 3D CNNs, they are limited to simple shapes, are relatively slow at inference, and have not yet been adapted for 3D semantic segmentation of medical data. Here, we propose Occupancy Networks for Semantic Segmentation (OSS-Nets) to accurately and memory-efficiently segment 3D medical data. We build upon the original O-Net with modifications for increased expressiveness leading to improved segmentation performance comparable to 3D CNNs, as well as modifications for faster inference. We leverage local observations to represent complex shapes and prior encoder predictions to expedite inference. We showcase OSS-Net's performance on 3D brain tumour and liver segmentation against a function space baseline (O-Net), a performance baseline (3D residual U-Net), and an efficiency baseline (2D residual U-Net). OSS-Net yields segmentation results similar to the performance baseline and superior to the function space and efficiency baselines. In terms of memory efficiency, OSS-Net consumes comparable amounts of memory as the function space baseline, somewhat more memory than the efficiency baseline and significantly less than the performance baseline. As such, OSS-Net enables memory-efficient and accurate 3D semantic segmentation that can scale to high resolutions.

研究の動機と目的

• 高解像度医用画像セグメンテーションにおける3D CNN の立方体メモリ成長問題に対処すること。
• 標準的なオカピシティネットワーク（O-Nets）の表現力の制限と遅い推論速度を克服すること。
• O-Nets のメモリ効率と3D CNN の性能を統合し、3次元セマンティックセグメンテーションを実現すること。
• メモリと推論時間を削減することで、コンsumer GPU 上での高解像度セグメンテーションを可能にすること。
• 脳腫瘍や肝臓病変のような複雑な解剖的構造を捉えることができる、スケーラブルで関数空間に基づく手法を開発すること。

提案手法

• OSS-Net は、複雑な3次元形状の表現能力を向上させるために、元の O-Net に局所観測を組み込む。
• 推論を高速化するため、事前エンコーダーの予測を初期状態として用い、実行時間を2倍短縮する。
• 1回の順伝播で複数の3次元位置での占有状態を評価する微分可能サンプリング戦略を採用する。
• マルチスケール特徴抽出機構により、ボリュメトリックデータ内の微細な解剖的詳細を捉える。
• 関数空間表現に適応した3D U-Net に類似たエンコーダ-デコーダーバッキングを活用する。
• 推論を最適化するため、事前予測に基づいて動的に位置をサンプリングし、精度を損なわず効率を向上させる。

実験結果

リサーチクエスチョン

• RQ1 オカピシティネットワークは、3次元医用画像における3D CNN と同等のセグメンテーション性能を達成できるように拡張可能か？
• RQ2 O-Nets の表現力をどのように向上させれば、腫瘍のような複雑な解剖的構造をモデル化できるか？
• RQ3 O-Nets の推論速度を顕著に向上させつつ、セグメンテーション精度を保持できるか？
• RQ4 関数空間3次元セグメンテーションにおいて、サンプリング位置とメモリ消費量の最適なトレードオフは何か？
• RQ5 LiTS 肝臓セグメンテーションデータセットのような限られた学習データを有する挑戦的なデータセットにおいて、OSS-Net はどのように性能を発揮するか？

主な発見

• BraTS 2020 脳腫瘍データセットにおいて、OSS-Net は3DリラクゼーションU-Net（性能ベースライン）と同等のセグメンテーション性能を達成し、Dice スコアは 0.878 を記録した。
• LiTS 肝臓セグメンテーションデータセットでは、O-Net ベースラインおよび2D U-Net 効率ベースラインを上回り、Dice スコア 0.852 を達成した。
• 推論速度は、元の O-Net アプローチと比較して2倍向上し、2^12 個を超えるサンプリング位置では実行時間が変化しなかった。
• 推論中のメモリ使用量は O-Net と同等（1.2GB）であり、3D U-Net ベースライン（6.5GB）の5倍以上も低減されており、コンsumer GPU へのデプロイが可能である。
• 元のアプローチと改善された推論アプローチの予測の不一致は 5.9658×10⁻⁴% にとどまり、性能の著しい低下は確認されなかった。
• OSS-Net は学習時にも 2.6GB のメモリを使用しており、3D U-Net ベースライン（6.5GB）と比較して顕著に低減されており、高い精度を維持している。