MONAIフレームワークを利用して、Amazon SageMakerで医用画像解析パイプラインを構築する

この記事は、“Build a medical image analysis pipeline on Amazon SageMaker using the MONAI framework” を翻訳したものです。

医用イメージングは、臨床医が疾患を特定し、より良い診断を行い、患者のアウトカムを劇的に改善できるようにする形で医療を変革しました。デジタル医用画像には、超音波、X線、コンピュータ断層撮影（CT）、磁気共鳴画像（MRI）、陽電子放射断層撮影（PET）、網膜写真、組織学スライド、皮膚鏡検査画像など、いくつかのモダリティ（撮影機器）があります。畳み込みネットワークなどのディープラーニングアルゴリズムの最近の開発は、診断と意思決定を補完し、臨床医に疾患に関する洞察を提供し、反復を自動化することで、医用画像のパターンを特定および分類するのに役立っています。タスクを実行し、高リスク状態の患者の迅速なトリアージを可能にします。

皮膚がん分類のためのDenseNet 畳み込みネットワークモデル

医用画像解析におけるディープラーニングには独自の課題があり、モデルのパフォーマンスを向上させるには、ドメイン固有のアプローチが必要です。たとえば、ディープラーニングでは、モデルをトレーニングするために大きなアノテーション付きデータセットが必要ですが、医用画像に使用できるパブリックデータセットは限られています。このためには、不均衡な観測値が存在するデータや不均衡なデータセットを扱うためのディープラーニング手法が必要です。医用画像も高解像度で、モデルのトレーニングやホスティングのために大量のストレージとコンピューティングを必要とするため、コストを最適化しながらこれらのインフラストラクチャ要件に対処するためのアーキテクチャが必要になります。さらに、ヘルスケア画像の研究のディープラーニングで用いられるソフトウェアツールキット（例：NiftyNet、DeepNeuro、Clara Train）には標準化や整合性はほとんどありません。そのため、研究者がベストプラクティスを定義したり、研究を再現するために協力することが難しくなっています。

2018 年、AWSは、研究者が公開データセットを共有できる「Registry of Open Data」をリリースしました。また、AWSは、研究者が精度の高いトレーニングセットを構築することを支援するデータラベリングワークフローを備えた「Amazon SageMaker Ground Truth」をリリースしました。Ground Truth は Amazon SageMaker スタックのコンポーネントの 1 つです。SageMakerは、研究者がサーバーを管理することなく、単一のインターフェイスで機械学習モデルを構築、トレーニング、デプロイできるフルマネージド型サービスです。また、分散環境で非常に大きなデータに対して実行できるように最適化された一般的な機械学習アルゴリズムを提供します。また、相互運用性を実現するために、独自のアルゴリズムとオープンソースフレームワークをサポートしています。SageMaker は、研究者が使用した分だけ課金され、高度にスケーラブルな環境でモデルを簡単にトレーニングおよびデプロイできます。

MEDICAL OPEN NETWORK FOR AI（MONAI）は 2020 年 4 月にリリースされた、PyTorchをベースとする、ヘルスケア画像におけるディープラーニングのための オープンソースフレームワークです。ウェブサイトに記載されているように、プロジェクトMONAIは「アカデミアと企業の研究者の間で、ヘルスケア画像に対するAIのベストプラクティスを開発・交換するために、AI研究者の包括的なコミュニティを設立する」取り組みです。MONAI の特徴は、多次元データの柔軟な前処理変換、既存のワークフローへの統合を容易にするためのポータブル API、ネットワーク、損失関数、評価指標のドメイン固有の実装などです。

この記事では、MONAI フレームワークを SageMaker マネージドサービスに統合する方法をご紹介します。また、不均衡なデータセットと画像変換を支援できる MONAI の前処理変換のサンプルコードも提供します。次に、画像分類のために DenseNet などの MONAI ニューラルネットワークアーキテクチャを呼び出すコードを確認し、SageMaker 内でモデルをトレーニングして提供するための PyTorch コードの構造を探ります。さらに、モデルトレーニングと推論のためのホスティングの両方のコンピューティングインフラストラクチャを起動および管理するための SageMaker API 呼び出しについても説明します。この例は、皮膚がん分類モデルに基づいており、ハーバード大学が発表したHAM10000 dermatoscopy skin cancer image dataset(皮膚鏡検査皮膚がん画像データセット)を使用しています。

前提条件

GitHub リポジトリでサンプルモデルコードを共有しました。モデルを自分のアカウントにデプロイしたい場合は、以下のことを確認してください。

• 管理者権限のある AWS アカウント
• GitHub アカウント
• Jupyter Notebook と Python の事前の経験をお勧めします

SageMaker で MONAI モデルをビルドする

SageMaker は、管理された Jupyter Notebookや Docker コンテナなど、研究者が慣れ親しんだツールを提供しています。ビルトインのアルゴリズムからモデルを構築したり、ビルトインコンテナを使用したスクリプトモード、独自のコンテナのデプロイ、 AWS Marketplace のコンテナの使用など、いくつかのオプションを使用して、独自のアルゴリズムを構築できます。

皮膚がんモデルでは、既存の PyTorch ビルトインコンテナを利用してスクリプトモードを使用し、カスタムアルゴリズム用の MONAI パッケージで拡張します。スクリプトモードでは、Jupyter Notebook とは別のファイルにトレーニングスクリプトを作成し、モデルトレーニングジョブの起動の際に使用します。皮膚がんの例では、monai_skin_cancer.py というスクリプトを作成しました。

トレーニングスクリプトと同じディレクトリに、 requirements.txt ファイルを作成して MONAI と、トレーニングをサポートするために PyTorch ビルトインコンテナにインストールする必要のある追加パッケージを含めます。

例:requirements.txt

monai==0.3.0
torchtext==0.6.0

学習スクリプトは、環境変数を通じてSageMakerのトレーニング環境に関するいくつかのプロパティにアクセスすることができます。例えば、ホストで利用可能なGPUの数、ハイパーパラメータ、モデルのアーティファクトのディレクトリなどです。モデルの学習のためにtrain関数を呼び出す下記の例のように、学習用のコードはmain guard(if name == '__main__':)の後に記述する必要があります。

例 monai_skin_cancer.py: PyTorch トレーニングスクリプトのentry point

#Main guard
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    
    # Container environment
    parser.add_argument('--hosts', type=list, default=json.loads(os.environ['SM_HOSTS']))
    parser.add_argument('--current-host', type=str, default=os.environ['SM_CURRENT_HOST'])
    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
    parser.add_argument('--data-dir', type=str, default=os.environ['SM_CHANNEL_TRAINING'])
    parser.add_argument('--num-gpus', type=int, default=os.environ['SM_NUM_GPUS'])

    train(parser.parse_args())
…

train 関数は、MONAI 変換を使用して皮膚がんトレーニングデータセットを取り込み、前処理と拡張を行います。MONAI は、医用画像の高次元性に特化した、辞書形式と配列形式の両方をサポートする変換機能を備えています。変換には、Crop & Pad、Intensity、IO、Post-processing、Spatial、Utilitiesなどのいくつかのカテゴリが含まれます。次の抜粋では、Compose クラスは一連の画像変換を連鎖させ、画像の単一のテンソルを返します。

例 monai_skin_cancer.py: MONAI 画像変換

#Load and transform skin cancer images
class SkinCancerDataset(Dataset):

    def __init__(self, image_files, labels, transforms):
        self.image_files = image_files
        self.labels = labels
        self.transforms = transforms

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, index):
        return self.transforms(self.image_files[index]), self.labels[index]

train_transforms = Compose([
        LoadPNG(image_only=True),
        AsChannelFirst(channel_dim=channel_dims_num),
        ScaleIntensity(),
        RandRotate(range_x=15, prob=0.5, keep_size=True),
        RandFlip(spatial_axis=0, prob=0.5),
        RandZoom(min_zoom=0.9, max_zoom=1.1, prob=0.5, keep_size=True),
        Resize(spatial_size=(64,64)),
        ToTensor()
    ])
    
dataset = SkinCancerDataset(image_files, labels, train_transforms)

train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) if is_distributed else None

train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=train_sampler is None,
                                       sampler=train_sampler, **kwargs)
…

#Instantiate model	
device = torch.device("cuda")
model = densenet121(
    spatial_dims=spatial_dims_num,
    in_channels=in_channels_num,
    out_channels=out_channels_num
).to(device)
loss_function = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), 1e-5)

#Train skin cancer model
for epoch in range(epoch_num):
    
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step += 1
        inputs, labels = batch_data[0].to(device), batch_data[1].to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        
        epoch_len = len(train_loader.dataset) // train_loader.batch_size        
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    
…

後のステップでトレーニングされた MONAI モデルを SageMaker にデプロイするには、トレーニングスクリプト内で、環境変数SM_MODEL_DIR によって定義されたパスへモデルを保存する必要があります。

例 monai_skin_cancer.py: モデルを保存する

path = os.path.join(model_dir, 'model.pth')
torch.save(model.state_dict(), path)

トレーニングジョブが完了すると、SageMaker はトレーニング済みモデルを圧縮し、ジョブで指定したAmazon Simple Storage Service (S3) バケットにアップロードします。次のセクションでは、トレーニングジョブを作成し、推論のためにモデルをデプロイする方法を説明します。

SageMaker 皮膚がんモデルのライフサイクル

MONAI モデルのトレーニングのスケールアップ

SageMaker は、各モデルのトレーニングにエフェメラルクラスターを使用します。エフェメラルクラスターは、トレーニングジョブが実行されている秒数だけ起動する専用インスタンスです。これにより、インスタンス上でマネージド Jupyter Notebook 内にモデルを構築し、独立したインスタンスでモデルをトレーニングして、パッケージのインストールやリソースの依存関係の競合を避け、コストを最適化できます。

次の例は、Jupyter Notebookで SageMaker Python SDK を使用してトレーニングジョブをインスタンス化し、前のセクションで説明した皮膚がんモデルのトレーニングスクリプトを実行する方法を示します。

例 monai_skin_cancer_notebook.ipynb: トレーニングスクリプトを実行するトレーニングジョブ

s3_input_train = sagemaker.s3_input(s3_data=skin_cancer_training)

estimator = PyTorch(entry_point='monai-skin-cancer-train.py',
                    role=role,
                    framework_version='1.5.0',
                    py_version=’py3’,
                    train_instance_count=1,
                    train_instance_type=ec2_instance_type,
                    hyperparameters={'epoch': '30'})

estimator.fit({'train': s3_input_train})

この例では fit メソッドを呼び出すと、train_instance_type パラメーターに基づくコンピュートサイズでトレーニングインスタンスを起動し、ビルトインの PyTorch ディープラーニングコンテナをインスタンス化し、requirements.txt で見つかった MONAI 依存パッケージをインストールし、Amazon S3から皮膚がんトレーニングデータセットがダウンロードされ、entry_pointパラメータで参照されるトレーニングスクリプトが実行されます。皮膚がんトレーニングデータセットのサイズが小さいため、トレーニングインスタンスのAmazon Elastic Block Store（EBS）ボリュームにダウンロードできました。しかし、大規模な医用画像データセットの場合、SageMaker パイプ 入力モードを使用して、トレーニングインスタンスにデータを直接ストリーミングできます。詳細については、「Amazon SageMaker アルゴリズムのパイプ入力モードを使用する」を参照してください。

また、SageMakerでは大規模データセットでのモデルの学習にPyTorch Estimatorを使用したマルチマシン分散学習をサポートしています。train_instance_countに1より大きい値を指定すると、SageMakerはトレーニングクラスター用に複数のインスタンスを起動し、各ホストでトレーニングスクリプトを実行します。このアプローチのデザインパターンと実装の詳細は、「AWS クラウドの GPU を使用した、スケーラブルなマルチノードの深層学習トレーニング」に記載されています。

SageMakerでは、メモリ、ディスク、GPUの使用率、トレーニングエラーなど、トレーニングジョブのモニタリングメトリクスが用意されており、モデルのトレーニングを可視化することができます。MONAIは、画像セグメンテーションのMean Diceなどの医用画像モデルの品質に焦点を当てた評価指標を提供します。SageMaker とMONAI の統合により、研究者は医用画像モデルのトレーニングを包括的に把握できます。

皮膚がんのモデルトレーニングでは、MONAIのメソッド compute_roc_auc を使用して、ROC（receiver operating characteristic）曲線の下の面積を測定しました。

例 monai_skin_cancer.py: ROC曲線の下の面積

for val_data in val_loader:
    val_images, val_labels = val_data[0].to(device), val_data[1].to(device)
    y_pred = torch.cat([y_pred, model(val_images)], dim=0)
    y = torch.cat([y, val_labels], dim=0)
    auc_metric = compute_roc_auc(y_pred, y, to_onehot_y=True, softmax=True)
    metric_values.append(auc_metric)
    acc_value = torch.eq(y_pred.argmax(dim=1), y)
    acc_metric = acc_value.sum().item() / len(acc_value)
    if auc_metric > best_metric:
         best_metric = auc_metric
         best_metric_epoch = epoch + 1
         torch.save(model.state_dict(), 'best_metric_model.pth')
         print('saved new best metric model')
    print(f"current epoch: {epoch + 1} current AUC: {auc_metric:.4f}"
          f" current accuracy: {acc_metric:.4f} best AUC: {best_metric:.4f}"
          f" at epoch: {best_metric_epoch}")
…

RESTful API エンドポイントに MONAI モデルをデプロイする

SageMaker は、推論サービスの可用性を高めるために、異なるアベイラビリティーゾーンにデプロイされる 1 つ以上のマネージド Amazon Elastic Compute (EC2) インスタンスで構成されるモデルエンドポイントを介してモデルをホストします。各 Amazon EC2 インスタンスには、各インスタンスが予測の応答を提供できるように、リクエストに応答するためのWebサーバーと、モデルのアーティファクトが含まれています。これらのインスタンスの状態は、リクエストを分散できるロードバランサーと、ピーク負荷時に新しいインスタンスを起動できる Auto Scaling によって維持されます。モデルエンドポイントは、estimator.deployを呼び出したときにSageMakerが特定のモデル用に作成するRESTful APIです。

AWS Lambdaをクライアントとして使用するSageMaker モデルエンドポイント

例 monai_skin_cancer_notebook.ipynb: マネージドエンドポイントにモデルをデプロイする

predictor = estimator.deploy(initial_instance_count=1, instance_type=ec2_instance_type)

作成したエンドポイントは、インスタンス上で PyTorch モデルサーバーを実行し、トレーニングスクリプトによって保存されたトレーニング済みモデルをロードします。サーバーは、MONAI 皮膚がんモデルのトレーニングスクリプトに含まれている model_fn を呼び出してモデルをロードします。

例 monai_skin_cancer.py: MONAI 皮膚がんモデルの読み込み

def model_fn(model_dir):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = densenet121(
        spatial_dims=spatial_dims_num,
        in_channels=in_channels_num,
        out_channels=out_channels_num
    )    
    with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
        model.load_state_dict(torch.load(f))
    return model.to(device)

SageMaker がトレーニング済みモデルをロードすると、predictor.predict を呼び出して推論リクエストする際に応答するモデルを提供します。皮膚がんモデルでは、患者の皮膚がんの画像を入力し、モデルはそのクラスと確率を予測しました。

例 monai_skin_cancer_notebook.ipynb: 皮膚がんのクラスを予測し、確率を示す

response = predictor.predict(data)
pred = torch.nn.functional.softmax(torch.tensor(response), dim=1)
top_p, top_class = torch.topk(pred, 1)
print("class: "+classes[top_class])
print("probablity: "+str(round(top_p.item(),2)))
Output> class: basal cell carcinoma, probability: 92 percent 

この例では、テストデータセット内の各皮膚がん画像のサイズが 2 メガバイト未満なので、リアルタイムの RESTful API エンドポイントを使用することができました。医用画像のサイズが大きい、または大量の画像を推測するユースケースは、推論にバッチ変換を使用することを検討する必要があります。バッチ変換ソリューションについては、「Amazon SageMaker の TensorFlow Serving を使用したバッチ推論の実行」で紹介されています。（訳注: またはAmazon SageMaker Asynchronous Inferenceをご検討ください。）

まとめ

SageMaker と MONAI の統合により、研究者は医用画像解析のベストプラクティスを定義し、共通の基盤として MONAI を使用して共同で研究することができます。SageMakerで提供される使い慣れたツールを、拡張性、安全性、コスト効率に優れたマネージド環境で活用することができます。この記事では、SageMaker の MONAI の基本的な機能をいくつか調べました。今後の研究としては、エンドツーエンドのワークフローのための教師あり学習エンジン、フェデレーテッドラーニング(連合学習)、イベント処理などが考えられます。

AWS の医療におけるイノベーションを加速する方法の詳細については、 aws.amazon.com/health をご覧ください。

翻訳は Solutions Architect 今井と窪田が担当しました。原文はこちらです。