[論文レビュー] Resource Optimized Neural Architecture Search for 3D Medical Image Segmentation

Woong Bae, Seungho Lee+4 more|arXiv (Cornell University)|Sep 2, 2019

Advanced Neural Network Applications参考文献 10被引用数 5

ひとこと要約

本論文は、1枚のRTX 2080Ti（10.8GBのGPUメモリ）を用い、1GBのデータセットでわずか1.39日間の学習で最先端の性能を達成する、リソース最適化されたニューラルアーキテクチャサーチ（NAS）フレームワーク、RONASMISを提案する。マクロサーチとパラメータ共有に注力し、再訓練を回避し、要素ごとの和算によるスキップ接続を採用することで、データ拡張、アンサンブル、テスト時拡張を一切用いずに、手動設計ネットワークおよび既存のNAS手法を上回る性能を達成する。

ABSTRACT

Neural Architecture Search (NAS), a framework which automates the task of designing neural networks, has recently been actively studied in the field of deep learning. However, there are only a few NAS methods suitable for 3D medical image segmentation. Medical 3D images are generally very large; thus it is difficult to apply previous NAS methods due to their GPU computational burden and long training time. We propose the resource-optimized neural architecture search method which can be applied to 3D medical segmentation tasks in a short training time (1.39 days for 1GB dataset) using a small amount of computation power (one RTX 2080Ti, 10.8GB GPU memory). Excellent performance can also be achieved without retraining(fine-tuning) which is essential in most NAS methods. These advantages can be achieved by using a reinforcement learning-based controller with parameter sharing and focusing on the optimal search space configuration of macro search rather than micro search. Our experiments demonstrate that the proposed NAS method outperforms manually designed networks with state-of-the-art performance in 3D medical image segmentation.

研究の動機と目的

• 3D医療画像セグメンテーションにおける既存のNAS手法の高い計算コストと長時間の学習を解決すること。
• 大規模かつアニューズロープな3D医療画像のNASにおいて、GPUメモリ使用量と学習時間を削減すること。
• 再訓練や大規模なデータ拡張を必要とせず、手動設計アーキテクチャを上回るNASフレームワークを開発すること。
• 限られた計算リソースの臨床現場でのNASの実用的導入を可能にすること。
• 戦略的なアーキテクチャ選択を伴うマクロサーチが、3D医療セグメンテーションにおいてマイクロサーチを上回ることを示すこと。

提案手法

• パラメータ共有を活用した強化学習ベースのコントローラーを用い、マクロサーチスペースを効率的に探索することで、学習コストとメモリ使用量を削減する。
• サーチスペースには、入力パッチサイズ、ダウンサンプリング操作、プーリングタイプ、拡張率、スキップ接続ポイントが含まれており、すべての選択が最小限のリソースで最大のパフォーマンスを発揮するように設計されている。
• パラメータとGPUメモリ消費量を削減するため、連結ではなく要素ごとの和算をスキップ接続に採用する。
• 深度方向畳み込みを標準3D畳み込みに置き換えることで、メモリ効率とパフォーマンスを向上させる。
• 学習中の子ネットワークの過学習を防ぐためにドロップパス正則化を適用し、バッチ正則化をインスタンス正則化に置き換えて学習の安定性を向上させる。
• コントローラーは150〜500エポックで学習され、子ネットワークは1エピソードあたり3エポックで学習され、報酬はバリデーションのDiceスコアに基づく。

実験結果

リサーチクエスチョン

• RQ1 限られたGPUリソースで動作する3D医療画像セグメンテーション向けのNASフレームワークを設計できるか？
• RQ2 戦略的なアーキテクチャ選択を伴うマクロサーチが、3D医療画像セグメンテーションタスクにおいてマイクロサーチを上回るか？
• RQ3 再訓練、データ拡張、アンサンブルを一切必要としないNAS手法が最先端のパフォーマンスを達成できるか？
• RQ4 パラメータ共有と要素ごとのスキップ接続は、パフォーマンスを維持しつつメモリ使用量をどれほど効果的に削減できるか？
• RQ5 事前学習モデルや複雑なデータパイプラインに依存せず、コントローラーに基づくNASアプローチが安定的かつ一貫性のあるアーキテクチャ探索を達成できるか？

主な発見

• 心臓セグメンテーションタスクにおいて、RONASMISは平均Diceスコア92.72を達成し、nnUnet（92.70）とSCNAS（89.99）を上回り、最先端の性能を示した。
• 前立腺タスクでは、テスト時拡張やアンサンブルを一切使用せず、Diceスコア75.71を達成し、nnUnet（74.54）とSCNAS（65.30）を上回った。
• 1枚のRTX 2080Ti（10.8GB GPUメモリ）を用い、1GBのデータセットで学習時間を1.39日間まで短縮し、従来のNAS手法と比べ顕著に低減した。
• コントローラーは80％以上の確率で同一または非常に類似したアーキテクチャを提案しており、エントロピーの低下と報酬の増加に伴い、安定的かつ一貫性のある探索行動を示した。
• データ拡張、アンサンブル、テスト時拡張、後処理を一切使用せず、優れたパフォーマンスを達成した。これは、大多数の最先端ベースラインとは対照的である。
• フレームワークは深さ方向畳み込みを標準3D畳み込みに置き換え、インスタンス正則化を採用することで、学習の安定性を向上させ、3D医療画像における効率的で安定したNASを実現した。