[論文レビュー] On the influence of Dice loss function in multi-class organ segmentation of abdominal CT using 3D fully convolutional networks

Chen Shen, Holger R. Roth+5 more|arXiv (Cornell University)|Jan 18, 2018

Radiomics and Machine Learning in Medical Imaging参考文献 6被引用数 31

ひとこと要約

本研究では、3D畳み込みニューラルネットワーク（FCN）を用いた腹部臓器の多クラスセグメンテーションにおいて、3つのDice損失重み付け戦略—均一、シンプル、および二乗—の影響を調査している。重み付けの選択と初期の学習率がセグメンテーション精度に顕著な影響を及ぼすことが示され、シンプル重み付け戦略が学習率0.01で最高の平均Diceスコア（81.0％）を達成した。

ABSTRACT

Deep learning-based methods achieved impressive results for the segmentation of medical images. With the development of 3D fully convolutional networks (FCNs), it has become feasible to produce improved results for multi-organ segmentation of 3D computed tomography (CT) images. The results of multi-organ segmentation using deep learning-based methods not only depend on the choice of networks architecture, but also strongly rely on the choice of loss function. In this paper, we present a discussion on the influence of Dice-based loss functions for multi-class organ segmentation using a dataset of abdominal CT volumes. We investigated three different types of weighting the Dice loss functions based on class label frequencies (uniform, simple and square) and evaluate their influence on segmentation accuracies. Furthermore, we compared the influence of different initial learning rates. We achieved average Dice scores of 81.3%, 59.5% and 31.7% for uniform, simple and square types of weighting when the learning rate is 0.001, and 78.2%, 81.0% and 58.5% for each weighting when the learning rate is 0.01. Our experiments indicated a strong relationship between class balancing weights and initial learning rate in training.

研究の動機と目的

• 3D腹部CTボリュームにおける多臓器セグメンテーション性能に、さまざまなDice損失重み付け戦略が与える影響を分析すること。
• Diceベースの損失関数を用いる際、初期学習率がトレーニング収束とセグメンテーション精度に与える影響を調査すること。
• ボクセル頻度が異なる多様な腹部臓器のセグメンテーションを向上させるために、最適な損失重み付けと学習率の組み合わせを特定すること。
• データ拡張なしの条件下で、3D U-Netモデルの一般化性能と安定性が異なる損失関数設定にどのように影響を受けるかを評価すること。
• 多クラス医療画像セグメンテーションタスクにおける3D FCNのハイパーパrameter選定に関する実証的指針を提供すること。

提案手法

• エンドツーエンドの3D多臓器セグメンテーションを実現するため、バックボーンネットワークとして3D U-Netアーキテクチャを採用した。
• 均一（Wu=1）、シンプル（Ws = N / (L|Rl| + ε)）、および二乗（Wq = N / (L|Rl|² + ε)）の3つのDice損失重み付け戦略を実装した。ここで|Rl|はクラスlに属するボクセル数を表す。
• 確率的勾配降下法を用いてトレーニングを行い、2つの初期学習率μ = 0.001およびμ = 0.01を用いた。
• 全CTボリュームから64×64×64のサブボリュームをランダムにクロップすることで、ミニバッチ学習とフルボリューム推論を可能にした。
• 8つの腹部臓器におけるセグメンテーション性能を、Dice類似係数（DSC）を用いて評価した。
• 重み付け式におけるゼロ除算を回避するため、ε = 1を適用した。

実験結果

リサーチクエスチョン

• RQ1 Dice損失関数の均一、シンプル、および二乗重み付けが、複数の腹部臓器におけるセグメンテーション精度にどのように影響を与えるか？
• RQ2 初期学習率（0.001対0.01）を変化させた場合、Diceベースの損失関数の収束性と性能にどのような影響が生じるか？
• RQ3 損失重み付けと学習率のどの組み合わせが、全臓器の平均Diceスコアを最大にするか？
• RQ4 クラス不均衡とボクセル頻度は、さまざまなDice損失重み付け戦略の有効性にどのように影響を及けるか？
• RQ5 トレーニングダイナミクス（例：損失曲線の挙動）は、重み付け戦略ごで顕著に異なるものとなり、収束パターンの違いを示唆するか？

主な発見

• 学習率0.001の条件下では、均一重み付け戦略が全8臓器で最高の平均Diceスコア（81.3％）を達成した。
• 学習率0.01の条件下では、シンプル重み付け戦略が他の戦略を上回り、平均Diceスコア81.0％を達成した。特に動脈、門脈、膵臓のセグメンテーション性能が顕著に向上した。
• 二乗重み付け戦略はμ = 0.01の条件下で顕著な改善を示し、平均Diceスコア58.5％を達成。特に肝臓、脾臓、胃のセグメンテーションが向上した。
• シンプル重み付け戦略は、学習率を引き上げることによる性能向上が最大で、μ = 0.001の59.5％からμ = 0.01の81.0％に上昇した。
• トレーニング曲線の分析から、均一重み付けモデルは10,000イテレーションで収束したが、シンプルおよび二乗戦略は10,000イテレーション以降も改善の余地があることが示された。
• データ拡張なしでも顕著な過学習は観察されず、トレーニングデータセットが十分に大きく多様であることが示唆された。