3D医学画像セグメンテーションにおけるモデル比較と改善戦略

プロジェクト概要

本プロジェクトは、多臓器CTデータセットを対象とした3D医学画像セグメンテーションを目的としています。アルゴリズム部分は主に二つのタスクから構成されています。

  1. 提供されたベースライン(nnU-Net)を基に、データ拡張、ハイパーパラメータチューニング、および粗分割から細分割へ(Coarse-to-Fine)の戦略を適用し、10分割交差検証による平均Dice係数を向上させる。
  2. PaddleSegフレームワークに内蔵されている6つのセグメンテーションモデル(VNet、UNETR、nnFormer、nnUNet、TransUNet、SwinUNet)を試行・比較し、詳細な実験プロセスと評価結果を提示する。

データセット概要

本データセットは、脾臓、右腎、左腎、胆のう、食道、肝臓、胃、大動脈、下大静脈、膵臓、膀胱の11の臓器のアノテーションを含む160の訓練データと40のテストデータで構成されています。

第一部分:ベースラインのアルゴリズム改善

本タスクの目的は、大会提供のベースラインを基に、より優れたアルゴリズムを設計し、より良い性能を達成することです。設計の考え方は、データセットの正規化後、回転、ガウスノイズ、ガウスぼかしといったデータ拡張手法を用いてモデルのロバスト性を向上させ、訓練時には10分割交差検証戦略を適用し、さらに最初の訓練結果から3Dデータの関心領域(ROI)を切り出し、より精密なセグメンテーション結果を得るために再訓練することです。

一、ベースライン:nnU-Netの紹介

nnU-Net(no-new-UNet)は、U-Netアーキテクチャの拡張と改良版です。U-Netは画像セグメンテーションに広く使用される畳み込みニューラルネットワーク(CNN)構造です。従来のU-Netと比較して、nnU-Netは、異なる医学画像セグメンテーションタスクに適応するための自己設定メカニズムを導入しており、ネットワークアーキテクチャを手動で調整および最適化する必要がありません。この自己設定機能により、nnU-Netは異なるデータセットやタスクに適用する際により柔軟で効率的になります。

nnU-NetフレームワークはU-Netの構造を大幅に変更するのではなく(ReLUをLeaky ReLUに、BNをInstance Normに変更するなど)、データの前処理(リサンプリングと正規化)、訓練(損失関数、オプティマイザ設定、データ拡張)、推論(パッチベース戦略、テスト時拡張、モデル統合など)、後処理(単一連結領域の強化など)に重点を置いています。

二、ベースラインの実行結果

公式提供のコードに基づき、ベースラインを10分割交差検証戦略で訓練・検証した結果は以下の通りです。

表1:各分割の全クラスの平均Dice係数

分割0 分割1 分割2 分割3 分割4 分割5 分割6 分割7 分割8 分割9 平均
Dice 0.8149 0.859 0.8831 0.8083 0.8046 0.8512 0.8497 0.8218 0.7486 0.8405 0.82817

表2:10分割の平均Dice係数(mDice)

背景 脾臓 右腎 左腎 胆のう 食道 肝臓 大動脈 下大静脈 膵臓 膀胱
Dice(オンライン) 0.99737853 0.91039857 0.92159785 0.90497482 0.64492576 0.68816753 0.94623484 0.77011192 0.91045195 0.81994966 0.69421586 0.72977343

三、アルゴリズムの改善

医学画像セグメンテーション分野では、深層学習モデルnnU-Netは基準となる優れたスタート地点を提供します。しかし、タスクの複雑さとデータの多様性が増すにつれて、データ拡張、10分割交差検証、ハイパーパラメータチューニング、およびCoarse-to-Fine戦略を組み合わせることで、Dice係数の値を向上させました。データ拡張はデータの多様性とモデルの汎化能力を向上させ、10分割交差検証は評価結果の信頼性を確保し、ハイパーパラメータチューニングは最適な設定を見つけ、Coarse-to-Fine戦略はセグメンテーション結果の精度と詳細な表現を強化します。これらの戦略は、nnU-Netモデルを効果的に最適化し、医学画像セグメンテーションの性能を大幅に向上させました。

1. データ拡張

本タスクでは、モデルの性能と汎化能力を向上させるため、いくつかのデータ拡張戦略を採用しました。主なデータ拡張手法には、回転、ガウスノイズ、ガウスぼかしが含まれます。**画像をランダムに回転(±15度)**し、**ガウスノイズ(10%)**を導入することで、画像のロバスト性を高め、現実のシナリオでのノイズ状況をシミュレートします。さらに、**ガウスぼかし処理(20%)**により、モデルはぼやけた画像に適応し、ぼやけた状況への対応能力を強化します。

以下の設定を追加しました:

transforms:
  - type: RandomRotation3D
  - type: GaussianNoiseTransform
  - type: GaussianBlurTransform

2. ハイパーパラメータチューニング

モデル改善の有効性を検証するため、ハイパーパラメータチューニングを行いました。バッチサイズ、学習率、オプティマイザの3つの重要なパラメータに注目し、その値を慎重に調整して比較実験を行い、モデルの収束を速め、性能を大幅に向上させる最適な設定を見つけました。

iter = 30000であればより良い結果が得られることは確かですが、時間やハードウェアの制約を考慮し、10分割で統一してイテレーション数を1.5万回に設定しました。

比較実験の結果は以下の通りです。iter=15000、stage=0で統一。

実験 オプティマイザ 学習率 バッチサイズ Dice
E1 sgd 0.01 2 0.7673
E2 adamw 0.01 2 0.8377
E3 adamw 0.02 2 0.83424

E1とE2の比較から、adamwはより速く収束し、同じiter数の場合、より優れたパフォーマンスを発揮することがわかります。E2とE3の比較から、学習率を大きくすると、元の学習率よりも効果が悪いことがわかります。

3. Coarse-to-Fine戦略

unet-cascadeのlowres+fullresがより良い性能を発揮する可能性があるものの、訓練時間が長すぎるため、依然としてunet級連構造stage=0を基礎モデルアーキテクチャとして採用しました。

本タスクでは、Coarse to Fine(粗分割から細分割へ)戦略を採用しました。これは、低解像度データを先に処理してから高解像度画像を処理するlower-fullres操作ではなく、まず粗分割を行い、モデルが大まかな目標領域を予測し、その領域から関心領域(ROI)を切り出し、ROIの範囲内でより精密なセグメンテーションを行い、詳細な特徴に注目してセグメンテーションの精度を向上させる戦略です。

具体的には、まず粗分割の結果からテストデータセットの初步的なセグメンテーション結果を得て、モデルが目標領域を大まかに予測し、これらの領域から関心領域(ROI)を切り出します。その後、ROIの範囲内でより精密なセグメンテーションを行い、詳細な特徴に注目してセグメンテーションの精度を向上させます。

具体的には、まず予測結果/ラベルから目標領域の最小矩形(rectangle)を計算します。次に、矩形の中心点と半径から、関心領域(ROI)の範囲を計算します。その後、元の画像、ラベル、および予測結果をROIの範囲に切り出し、関心のある部分を得ます。

擬似コードは以下の通りです。

# 変数を宣言
rectangle_indices <- 非ゼロインデックス(target['arr'])
rectangle <- [
    [min(rectangle_indices[0]), max(rectangle_indices[0])],
    [min(rectangle_indices[1]), max(rectangle_indices[1])],
    [min(rectangle_indices[2]), max(rectangle_indices[2])],
]
center <- [(a + b) / 2 for a, b in rectangle]
radius <- [(b - a) / 2 * zoom for a, b in rectangle]
start_position <- [r[0] for r in rectangle]

# ループ処理
FOR EACH key IN d DO
    IF d[key]['arr'] IS None THEN
        CONTINUE  # 空の値がある場合は、このループをスキップして次のキーを処理
    ENDIF

    # 関心領域(ROI)を切り抜く
    d[key]['arr'] <- (
        d[key]['arr'][
            rectangle[0][0] : rectangle[0][1],
            rectangle[1][0] : rectangle[1][1],
            rectangle[2][0] : rectangle[2][1],
        ]
    )
ENDFOR

4. ベースラインとCoarse-to-fineの結果比較

表1:改善後の各分割の全クラスの平均Dice係数
分割0 分割1 分割2 分割3 分割4 分割5 分割6 分割7 分割8 分割9 平均
平均Dice 0.8603 0.8591 0.8494 0.8541 0.8431 0.8516 0.8539 0.8256 0.8236 0.8762 0.850
表2:改善後の10分割の平均Dice係数(mDice)
背景 脾臓 右腎 左腎 胆のう 食道 肝臓 大動脈 下大静脈 膵臓 膀胱
mDice(オンライン) 0.99175904 0.91310376 0.92185367 0.91565554 0.66577194 0.75998226 0.94889251 0.79217957 0.93097676 0.85524831 0.74284509 0.75815549

上記の図から、iter=15000のとき、改善されたアルゴリズムモデルの結果がより安定していることがわかります。今後、条件が許せば、イテレーション数を増やした後の効果を探索する予定です。

四、実行フローパイプライン

以下はプロジェクトの全体的な実行フローです。実行結果の表示を減らすため、ここではコード実行手順のみを提示します。

# ステップ1: PaddleSegリポジトリをクローン
%cd ~/
!git clone --branch develop --depth 1 https://github.com/PaddlePaddle/PaddleSeg.git
%cd ~/PaddleSeg

# ステップ2: データを~/PaddleSeg/contrib/MedicalSeg/data/raw_dataに解凍
!mkdir ~/PaddleSeg/contrib/MedicalSeg/data
!unzip -oq ~/data/data204195/base_train.zip -d ~/PaddleSeg/contrib/MedicalSeg/data/raw_data

# ステップ3: 元のデータセット圧縮ファイルを削除してメモリを節約
!rm -rf ~/data/data204195/base_train.zip
# データセットに役に立たないフォルダがある場合は削除しないとエラーが発生する
!rm -rf ~/PaddleSeg/contrib/MedicalSeg/data/raw_data/.ipynb_checkpoints

# ステップ4: 必要な依存パッケージをインストール
%cd ~/PaddleSeg/contrib/MedicalSeg/
!pip install 'protobuf<=3.20.0' -r requirements.txt # バージョン競合を避ける
!pip install medpy
# ステップ5: ymlファイルにデータ拡張パラメータを追加した後、10分割交差検証を実行し、粗分割結果を得る。train.shに書き込んで直接実行することも可能だが、ここでは10分割訓練を示すため、コマンドラインの簡潔さは考慮していない。
!python train.py --config ~/configs/coarse/nnunet_1_fold2.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold2  --save_interval 1500 --use_vdl --do_eval    --resume_model output/cascade_lowres/fold2_adamw/iter_4500
!python train.py --config ~/configs/coarse/nnunet_1_fold0.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold0 --save_interval 1500 --use_vdl --do_eval    --resume_model output/cascade_lowres/fold0_adamw/iter_12000
!python train.py --config ~/configs/coarse/nnunet_1_fold1.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold1  --save_interval 1500 --use_vdl --do_eval    
!python train.py --config ~/configs/coarse/nnunet_1_fold3.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold3  --save_interval 1500 --use_vdl --do_eval    
!python train.py --config ~/configs/coarse/nnunet_1_fold4.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold4  --save_interval 1500 --use_vdl --do_eval  
!python train.py --config ~/configs/coarse/nnunet_1_fold5.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold5  --save_interval 1500 --use_vdl --do_eval  
!python train.py --config ~/configs/coarse/nnunet_1_fold6.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold6  --save_interval 1500 --use_vdl --do_eval   
!python train.py --config ~/configs/coarse/nnunet_1_fold7.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold7 --save_interval 1500 --use_vdl --do_eval   
!python train.py --config ~/configs/coarse/nnunet_1_fold8.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold8  --save_interval 1500 --use_vdl --do_eval   
!python train.py --config ~/configs/coarse/nnunet_1_fold9.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold9 --save_interval 1500 --use_vdl --do_eval    
# ステップ6: ROIを取得
%cd ~/PaddleSeg/contrib/MedicalSeg/
!python get_ROI.py --roi_data './data/02-roi/raw_data' --predict_data '../../../submit'
# データセットに役に立たないフォルダがある場合は削除しないとエラーが発生する
!rm -rf ~/PaddleSeg/contrib/MedicalSeg/data/02-roi/raw_data/.ipynb_checkpoints
# ステップ7:Coarse-to-fineの細分割部分。この部分のymlファイルは~/configs/fineに保存されています。
%cd ~/PaddleSeg/contrib/MedicalSeg/
!python train.py --config ~/configs/fine/nnunet_2_fold0.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold0 --save_interval 1500 --use_vdl --do_eval     --resume_model output/fine/fold0/iter_3000

!python train.py --config ~/configs/fine/nnunet_2_fold1.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold1 --save_interval 1500 --use_vdl --do_eval    --resume_model output/fine/fold1/iter_3000

!python train.py --config ~/configs/fine/nnunet_2_fold2.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold2 --save_interval 1500 --use_vdl --do_eval     #--resume_model output/fine/fold0/iter_3000

!python train.py --config ~/configs/fine/nnunet_2_fold3.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold3 --save_interval 1500 --use_vdl --do_eval     #--resume_model output/fine/fold0/iter_3000

!python train.py --config ~/configs/fine/nnunet_2_fold4.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold4 --save_interval 1500 --use_vdl --do_eval     #--resume_model output/fine/fold0/iter_3000

!python train.py --config ~/configs/fine/nnunet_2_fold5.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold5 --save_interval 1500 --use_vdl --do_eval     #--resume_model output/fine/fold0/iter_3000

!python train.py --config ~/configs/fine/nnunet_2_fold6.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold6 --save_interval 1500 --use_vdl --do_eval     #--resume_model output/fine/fold0/iter_3000

!python train.py --config ~/configs/fine/nnunet_2_fold7.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold7 --save_interval 1500 --use_vdl --do_eval     #--resume_model output/fine/fold0/iter_3000

!python train.py --config ~/configs/fine/nnunet_2_fold8.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold8 --save_interval 1500 --use_vdl --do_eval     #--resume_model output/fine/fold0/iter_3000

!python train.py --config ~/configs/fine/nnunet_2_fold9.yml \
        --log_iters 20 --precision fp16 --nnunet --save_dir output/fine/fold9 --save_interval 1500 --use_vdl --do_eval     #--resume_model output/fine/fold0/iter_3000

第二部分:PaddleSeg内蔵モデルの分析

一、TransUNet

1. ネットワークの紹介

TransUNetは、シーケンスからシーケンスの予測の観点から自己注意メカニズムを導入した最初の医学画像セグメンテーションフレームワークです。Transformerがもたらす特徴解像度の損失を補うため、TransUNetは混合CNN-Transformerアーキテクチャを採用し、CNN特徴の詳細な高解像度空間情報とTransformerエンコーダのグローバルな文脈情報を利用します。U型アーキテクチャ設計のインスピレーションを受け、Transformerエンコーダの自己注意特徴はアップサンプリングされ、エンコーダパスからスキップされた異なる高解像度のCNN特徴と結合され、正確な位置決めを実現します。実証結果は、以前のCNNベースの自己注意方法と比較して、Transformerベースのアーキテクチャが自己注意をより効果的に利用していることを示しています。

2. データ結果の分析

30,000回のイテレーションで訓練した後、モデルは検証セットで非常に良い結果を達成し、mDiceは0.8589でした。各臓器のDice係数は以下の通りです。

[EVAL] Class dice: 
[0.9977 0.9291 0.9212 0.9259 0.7739 0.6833 0.9631 0.7501 0.9169 0.8454 0.7745 0.8261]
[EVAL] The model with the best validation mDice (0.8589) was saved at iter 30000.

3. パイプライン

完全なnotebookは以下のリンクでご覧いただけます(リンクは削除済み)。ここではコードのみを提示します。

mkdir -p ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training
unzip -oq ~/data/data204195/base_train.zip -d ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training
mv ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training/imagesTr ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training/img
mv ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training/labelsTr ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training/label

その後、データの前処理を行い、まず元のデータ前処理スクリプトを修正したものに置き換えます。

# prepare_abdomen.pyの修正例
label_map = {
    str(i):i for i in range(12)
}
...
def generate_txt(self, train_split=0.9):
    ...
def train_val_split(self, train_split=0.9):
    ...
dataset_description="3D Smart Med",
...
cd ~/PaddleSeg/contrib/MedicalSeg
cp -f ~/prepare_abdomen.py tools/
python tools/prepare_abdomen.py

修正したモデル設定ファイルを使用して訓練を実行します。

python train.py --config ~/transunet_abdomen_224_224_1_14k_1e-2.yml --do_eval --save_interval 1000 --has_dataset_json False --is_save_data False --num_workers 4 --log_iters 10 --use_vdl

4. まとめ

TransUNetの著者は、Transformerの利点を最大限に活用するためにTransUNetを提案しました。それは画像特徴をシーケンスとしてエンコードし、強力なグローバルな文脈情報を獲得し、U型混合アーキテクチャ設計を採用することで、低レベルのCNN特徴を効果的に利用します。全畳み込みベースの医学画像セグメンテーション方法の代替フレームワークとして、TransUNetは、CNNベースの自己注意方法を含む複数の競合手法の中で優れたパフォーマンスを発揮します。そのパフォーマンスは優れており、私たちの訓練はまだその完全な実力を示していませんが、グラフから見ると、mDiceはさらに向上する可能性が高いです。

二、SwinUNet

1. ネットワークの紹介

近年、畳み込みニューラルネットワーク(CNN)は医学画像分析において大きな進歩を遂げています。特にU型アーキテクチャとスキップ接続に基づく深層ニューラルネットワークは、多种の医学画像タスクに広く応用されています。しかし、CNNは優れたパフォーマンスを発揮するにもかかわらず、畳み込み操作の局所性により、グローバルおよび長距離のセマンティック情報の相互作用を十分に学習できないという欠点があります。Swin-Unetは、Unet構造に類似した医学画像セグメンテーションのための純粋なTransformerネットワークです。

Swin-Unetは、トークン化された画像ブロックをTransformerベースのU型エンコーダ-デコーダ構造に入力し、スキップ接続を使用して局所的およびグローバルなセマンティック特徴学習を行います。具体的には、階層的なSwin Transformerを使用し、スライディングウィンドウをエンコーダとして文脈特徴を抽出します。対称的なSwin Transformerベースのデコーダは、patch expanding層と組み合わせてアップサンプリング操作を実行し、特徴マップの空間解像度を回復します。

Swin-Unetは複数のデータセットでTrans-UNetよりも良い結果を得ていますが、私たちのデータセットではTrans-UNetよりも劣るパフォーマンスを示しました。

2. データ結果の分析

モデルを30,000回のイテレーションで訓練した後、最良のモデルは検証セットでmDiceが0.8378を達成し、Trans-UNetよりわずかに劣ります。

[EVAL] #Images: 16, Dice: 0.8378, Loss: 0.104690
[EVAL] Class dice: 
[0.9964 0.9219 0.905  0.9059 0.7626 0.6477 0.9577 0.7332 0.8827 0.8139 0.7326 0.7938]
[EVAL] The model with the best validation mDice (0.8378) was saved at iter 29000.

3. パイプライン

完全なnotebookは以下のリンクでご覧いただけます(リンクは削除済み)。ここではコードのみを提示します。

mkdir -p ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training
unzip -oq ~/data/data204195/base_train.zip -d ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training
mv ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training/imagesTr ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training/img
mv ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training/labelsTr ~/PaddleSeg/contrib/MedicalSeg/data/abdomen/abdomen_raw/RawData/RawData/Training/label

データの前処理はTrans-UNetと同じ修正済みスクリプトを使用します。

cd ~/PaddleSeg/contrib/MedicalSeg
cp -f ~/prepare_abdomen.py tools/
python tools/prepare_abdomen.py

修正したモデル設定ファイルを使用して訓練を実行します。

_base_: 'abdomen.yml'

batch_size: 24
iters: 30000

model:
  type: SwinUNet
  backbone:
    type: SwinTransformer_tinier_patch4_window7_224
  num_classes: 12
  pretrained: https://paddleseg.bj.bcebos.com/paddleseg3d/synapse/abdomen/swinunet_abdomen_224_224_1_14k_5e-2/swinunet_pretrained.zip

optimizer:
  type: sgd
  momentum: 0.9
  weight_decay: 1.0e-4

lr_scheduler:
  type: PolynomialDecay
  learning_rate: 0.05
  end_lr: 0
  power: 0.9

loss:
  types:
    - type: MixedLoss
      losses:
        - type: CrossEntropyLoss
        - type: DiceLoss
      coef: [0.4, 0.6]
  coef: [1]

export:
  transforms:
    - type: Resize3D
      size: [ 1 ,224, 224 ]
  inference_helper:
    type: InferenceHelper2D
!python train.py --config ~/swinunet_abdomen_224_224_1_14k_5e-2.yml --do_eval --save_interval 1000 --has_dataset_json False --is_save_data False --num_workers 4 --log_iters 10 --use_vdl # --resume_model output/iter_14000

4. まとめ

Swin-UNetは名前にUNetが含まれていますが、Trans-UNetとは異なり、純粋なTransformerネットワーク構造であり、U型です。非常に新しいもので、論文の説明ではTrans-UNetよりも効果が良いはずですが、私たちのデータセットではそうではありません。もちろん、これは使用したパラメータにも関係がある可能性があります。グラフから見ると、訓練はほぼ収束していますが、さらなる訓練でより高いmDiceを得られる可能性も排除できません。

三、nnUNet

1. ネットワークの紹介

nnU-Net(no-new-Net)は3つのネットワークを提案しました:2D U-Net、3D U-Net、および級連Unet。これらのU-Netは元のU-Netに対して非常に小さな変更しか行っておらず、様々な拡張プラグイン(残差接続、Dense接続、様々な注意メカニズムなど)は使用していません。著者は、より多くの向上は実際にはデータを理解し、現在のデータセットに適した訓練プロセスを構築することにあると考えています。これには、データセットの前処理、モデル構造の自己適応調整、訓練戦略、および推論戦略が含まれます。著者はこのフレームワークを現在のSOTA(State-Of-The-Art)方法と比較し、その方法は手動でハイパーパラメータを調整する必要がないことを示しました。最終的にnnUNetは最高の平均Diceを獲得しました。

2. データ分析

...

四、VNet

1. ネットワークの紹介

V-NetはU-Netの改良ネットワークであり、構造的にはU-Netと非常に類似しています。主に医学画像セグメンテーションに使用されるネットワークです。U-Netと比較して、V-Netの主な改良点は以下の通りです:

  • V-Netは3D画像向けに設計されたモデルであり、U-Netは2D画像向けに設計されたモデルです。nnUNetと同様に、VNetはConv 3Dを使用してConv 2Dを置き換えています。
  • V-NetはResNetを模倣して残差接続を導入し、ReLU関数の代わりにPReLU活性化関数を使用して、ネットワークの表現能力と収束速度を強化します。
  • V-Netは、下サンプリングのプーリング(MaxPooling)層の代わりに2x2 Conv層を使用し、情報抽出と抽象化の能力をさらに強化します。

2. データ結果の分析

SGDオプティマイザを使用し、初期学習率を0.05とし、多項式減衰を採用して30,000回のイテレーションを行った後、V-Netの訓練曲線は以下の通りです。

各クラスのDice値は以下の通りです。

背景 脾臓 右腎 左腎 胆のう 食道 肝臓 大動脈 下大静脈 膵臓 膀胱
Dice(オンライン) 0.9945 0.8183 0.855 0.8588 0.4388 0.2249 0.9441 0.7753 0.8678 0.7536 0.6404 0.7198

結果から、V-NetのDiceは最終的に0.74前後でしか達成できず、nnUNetベースラインの0.8には及びません。V-Netの論文と比較し、nnUNetと比較すると、V-Netの焦点は依然としてネットワーク構造の改良にあります。これらの改良はネットワークの効果をU-Netよりわずかに優れたものにするかもしれませんが、依然として二段階訓練のnnUNetには及びません。

五、UNETR

1. ネットワークの紹介

UNETRは、3D医学画像セグメンテーションのための新しいアーキテクチャです。Transformerをエンコーダとして使用し、入力ボリュームのシーケンス表現を学習し、効果的にグローバルなマルチスケール情報をキャプチャし、成功した「U型」ネットワーク設計にも従います。Transformerエンコーダは、異なる解像度のスキップ接続を介して直接デコーダに接続され、最終的なセマンティックセグメンテーション出力を計算します。Multi Atlas Labeling Beyond The Cranial Vault (BTCV)データセットで多臓器セグメンテーションを実行し、Medical Segmentation Decathlon (MSD)データセットで脳腫瘍および脾臓セグメンテーションタスクの性能を検証しました。BTCVランキングで新しい最先端のパフォーマンスを示しました。

2. データ結果の分析

SGDオプティマイザを使用し、初期学習率を0.05とし、30,000回のitersで訓練した後、diceとlossの変化は以下の通りです。

最良モデルの検証後、各クラスのdice値は以下の通りです。

背景 脾臓 右腎 左腎 胆のう 食道 肝臓 大動脈 下大静脈 膵臓 膀胱
Dice(オンライン) 0.9911 0.7708 0.6734 0.6834 0.3001 0.1243 0.8765 0.578 0.7674 0.532 0.5385 0.6123

検証セット上の最良のdiceは約15,000 iters付近で停滞し、最良のdiceは0.62にしか達せず、Transformerを導入していないVNetにも及びません。

15,000 iters以降、訓練セット上のloss(交差エントロピー損失とdice損失を含む)は減少して最適化されていますが、検証セット上のdiceも不断に減少しており、過学習の兆候を示しています。

訓練リソースの制約により、パラメータをさらに調整することはありませんでしたが、より良いハイパーパラメータが存在する可能性があります。

六、nnFormer

1. ネットワークの紹介

自然言語処理分野では、Transformerはデフォルトのモデルとなっていますが、医学画像処理分野ではまだあまり注目されていません。Transformerが長い依存関係を利用できるため、伝統的な畳み込みニューラルネットワークが持つ空間帰納バイアスの固有の欠点を克服するのに役立つと期待されています。しかし、最近提案されたほとんどのTransformerベースのセグメンテーション方法は、Transformerを単に補助モジュールとして扱い、グローバルな文脈を畳み込み表現にエンコードするために使用しています。nnFormerは、ボリューム医学画像セグメンテーションのための3D Transformerです。

nnFormerは、交差した畳み込みと自己注意操作の組み合わせを利用するだけでなく、ローカルおよびグローバルなボリュームに基づく自己注意メカニズムを導入し、ボリューム表現を学習します。さらに、nnFormerは、U-Netに類似したアーキテクチャのスキップ接続における伝統的な接続/和算操作の代わりにスキップ注意(skip attention)を使用することを提案しています。実験は、以前のTransformerベースの方法よりも優れたパフォーマンスを示しています。

2. パイプライン

完全なnotebookは以下のリンクでご覧いただけます(リンクは削除済み)。ここではコードのみを提示します。

mkdir -p ~/PaddleSeg/contrib/MedicalSeg/data/ACDCDataset/training
unzip -oq ~/data/data204195/base_train.zip -d ~/PaddleSeg/contrib/MedicalSeg/data/ACDCDataset/training
rm -rf ~/PaddleSeg/contrib/MedicalSeg/data/ACDCDataset/training/.ipynb_checkpoints

元のデータ前処理スクリプトを私たちのデータセットに合わせて修正します。

# prepare_acdc.pyの修正例
def generate_txt(self, split=0.1):
    ...
dataset_description="3D Smart Med",
...
cd ~/PaddleSeg/contrib/MedicalSeg/
cp ~/prepare_data.py ~/PaddleSeg/contrib/MedicalSeg/tools
python tools/prepare_data.py

修正したnnFormerモデルに合わせて、設定ファイルの一部のパラメータを調整しました。

_base_: 'acdc_4e-4.yml'

model:
  type: nnFormer
  crop_size: [64,160,160]
  embedding_dim: 192
  input_channels: 1
  num_classes: 12
  depths: [2, 2, 2, 2]
  num_heads: [3, 6, 12, 24]
  patch_size: [2, 4, 4]
  window_size: [4,8,8,4]
  deep_supervision: True
cp ~/nnFormerNew.py medicalseg/models/nnformer.py
python train.py --config ~/nnformer_acdc_160_160_14_250k.yml --save_interval 2000 --num_workers 4 --do_eval --log_iters 250 --is_save_data False --has_dataset_json False

3. まとめ

nnFormerは非常に新しい3D Transformerです。残念ながら、十分な計算リソースがないため、このモデルを訓練することができませんでした。しかし、データ前処理スクリプト、nnFormerのコード、設定ファイルを修正し、このネットワークが私たちのデータセットで訓練できるようにしました。このモデルを訓練する興味のある読者は、私たちが使用したパラメータ値は参考程度であることに注意してください。

第三部分:まとめ

私たちのグループはPaddlePaddleのMedicalSegをツールとして使用し、NNUNetの改良を実現し、同じデータセット上のMedicalSegの異なるネットワークのパフォーマンスを比較しました。各ネットワークのiters = 15000のときの検証セット上のdiceは以下の通りです。

検証セットの結果から、nnUNetは、Trans-UNetなどのTransformerベースのネットワークと同様に、検証セットの効果を達成できることがわかります(diceはいずれも約0.85に達します)。nnUNetは軽量かつ実用的な特徴を兼ね備え、医学画像セグメンテーションにおいて強いパフォーマンスと汎化能力を持ち、応用の将来性が広いです。

タグ: 医学画像セグメンテーション PaddleSeg 3D CNN Transformer nnUNet

6月27日 17:21 投稿