AWS で実行する NVIDIA Clara Train アプリケーション開発フレームワークを使った医療画像のモデルトレーニングと AI 支援アノテーションの迅速化

2020 年 5 月、AWS は NVIDIA Clara Train アプリケーションフレームワークを使った医療画像モデルの開発環境を AWS クラウドにデプロイするために使用できる AWS クイックスタートをリリースしました。Philips および Cerner などの多数のヘルスケアおよびライフサイエンス関連のお客様が、機密性の高いヘルスケアワークロードに関して AWS を信頼しておられます。セキュアでスケーラブルなクラウドインフラストラクチャは、データサイエンティスト、アーキテクト、および医療研究者が、医療コミュニティを支援するための機械学習 (ML) のワークフローとパイプラインの構築に集中できるようにします。

FDA に認可された AI は、すでに世界中の病院で医師たちが時宜を得た高品質のケアを提供するために役立っています。商用ソリューションには AWS で本番環境の医療画像ワークロードを確立しているものがあり、これには GE ヘルスケア、Zebra、Arterys、および Heartflow によって開発されたソリューションが含まれます。これらの成功にもかかわらず、ヘルスケアにおける革命はまだ初期の段階にあり、データと ML のパイプラインが臨床医による診断、ワークフローの最適化、医療機器の強化、および臨床品質の向上を補助する機会が今後ますます増えるでしょう。

以下のセクションでは、NVIDIA Clara Train アプリケーションフレームワークを紹介してから、CT スキャンにおける脾臓のセマンティックセグメンテーションに関する実際のユースケースに対応するためのデプロイメントについて説明します。このクイックスタートは、Amazon Elastic Container Service (ECS) と、NVIDIA V100 Tensor Core GPU 搭載の Amazon EC2 p3.2xlarge インスタンスタイプを使って Clara Train SDK をデプロイします。

このブログ記事はデプロイメント手順を説明するクイックスタートの実践ガイドで、AWS Cloud9 統合開発環境 (IDE) のセットアップ、NVIDIA NGC (GPU 最適化済みソフトウェアのハブ) からの事前訓練済みモデルのロード、および Clara Train サーバーの Slicer (オープンソースの医療画像インフォマティックスツール) への接続を行います。また、Clara Train コマンドラインツールを使ってモデルトレーニングジョブを実行する方法も説明します。

NVIDIA Clara の概要

NVIDIA Clara は、医療画像分析とゲノミクスプロファイリングの両方を迅速化するためのツールのセットです。この記事では、特に医療画像に焦点を当てます。医療画像のための Clara アプリケーションフレームワークは、特定の目的を念頭に置いたワークロードに分割することができるコンテナ化されたソリューションです。そのひとつ目は Clara Train であり、これは AI 支援アノテーション (AIAA)、および AutoML、転移学習、および連合学習といった手法を使った AI モデルのトレーニングのための機能を提供します。ふたつ目は Clara Deploy で、Triton 推論サーバーで演算子とパイプラインを使用することによる推論のためのフレームワークを提供します。今回のクイックスタートは、Clara Train アプリケーションフレームワークをデプロイします。

AWS での Clara Train の使用は、以下のメリットを提供します。

大規模な医療画像データセットを耐久性に優れた方法で保存し、チーム間におけるセキュアなコラボレーションを促進するために使用できる、Amazon S3 などの AWS ストレージサービスへの素早いアクセス。
恒久的なコンピューティングインフラストラクチャの最小限化。デプロイメントはすべてエフェメラルで、必要に応じた GPU コンピューティングへのスケーラブルなアクセスを提供し、料金は従量制です。
Amazon ECS は完全マネージド型のコンテナオーケストレーションサービスであるため、Clara Train Docker コンテナと基盤となる GPU コンピューティングも管理されます。サービスレベルのヘルスチェックは、Clara Train サービスの稼働時間を確保し、障害が発生した場合には、自動的に 2 番目の AWS アベイラビリティーゾーンにフェイルオーバーされます。
高可用性アーキテクチャにデプロイされた Clara Train の AIAA サービスは、医療画像データセットのアノテーションに関連する労力とコストを大幅に削減できます。
AutoML、連合学習、および転移学習といった Clara Train の高度な機能を導入することで、高機能の ML パイプラインの開発を迅速化できます。

つまり、Amazon ECS、Amazon Elastic File System (EFS)、および Amazon S3 などのマネージドサービスとの Clara Train の活用は、基盤となるプラットフォームの管理に関連する「Undifferentiated Heavy Lifting」(差別化につながらない重労働) ではなく、組織の研究または製品の差別化につながる AI 開発に焦点を当てる上で役立ちます。

前提条件

デプロイメントガイドの手順に従って、NVIDIA Clara Train SDK クイックスタートをデプロイします。デフォルトで、クイックスタートは AWS Certificate Manager で管理されるお客様の証明書を使って、環境全体におけるネットワーク接続のエンドツーエンドの暗号化を行います。モデルトレーニングと連合学習重視のデプロイメントには、暗号化を必須とするデフォルト設定を使用してください。

Slicer との統合を希望する場合は、デフォルト設定 (およびベストプラクティス) を上書きし、ロードバランサーとアノテーションサーバーへの暗号化されていない HTTP 接続を許可する必要があります。この安全性が低い設定でアーキテクチャをデプロイするには UseHTTPS を HTTP に設定し、匿名化された医療画像データでの設定以外では使用しないように注意してください。

デプロイメントが完了したら、AWS CloudFormation の [出力] セクションに Clara AIAA API と Amazon EFS DNS の名前が表示されます。

ClaraServiceUrl 出力をブラウザに貼り付けます。Clara AIAA サーバーが提供する API ドキュメントページで、デプロイメントが正常に行われたことを確認します。

Medical Segmentation Decathalon から、ローカルマシンに Task09_Spleen.tar データセットをダウンロードします。アーカイブを S3 バケットにコピーし、アーカイブの内容をローカルマシンのディレクトリに抽出します。

IDE および環境設定

SDK 機能を試してみるために、事前訓練済みモデルとトレーニングデータの両方をセットアップします。この設定は、クイックスタートに含まれるスクリプトによって、部分的に自動化されています。setup スクリプトの実行と、モデル設定ファイルの編集には、AWS Cloud9 IDE 環境を使用します。AWS Cloud9 EC2 インスタンスは、クイックスタートでデプロイされた Amazon EFS ファイルシステムをマウントし、Clara Train コンテナが使用するデータと設定へのアクセスを提供します。

AWS マネジメントコンソールから、新しい AWS Cloud9 環境を立ち上げます。ネットワーク設定 (詳細) 以外はデフォルト設定を受け入れます。AWS Cloud9 EC2 インスタンスを、ECS クラスターが含まれるものと同じ VPC に起動します。サブネットには、同じ VPC のパブリックサブネットを選択します (踏み台ホストが含まれるサブネットなど)。

AWS Cloud9 IDE の下側のペインにある bash シェルで、クイックスタートリポジトリをローカルディレクトリにクローンします。

git clone https://github.com/aws-quickstart/quickstart-nvidia-imaging-clara-train.git

AWS Cloud9 エディタを使って、クローンしたリポジトリの samples サブディレクトリにある clara_bootstrap.sh スクリプトを開きます。CLOUD9SG を AWS Cloud9 EC2 インスタンスが使用するセキュリティグループ ID に設定します。EFSSG を、クイックスタートによってデプロイされた Amazon EFS マウントターゲットが使用するセキュリティグループ ID に設定します。AIAA および EFSDNS はそれぞれ、AWS CloudFormation スタック出力の ClaraServiceUrl と ElasticFileSystemDnsName に設定します。DATABUCKET を、データアーカイブ例の Task09_Spleen.tar をアップロードした S3 バケットに設定します。

ファイルを保存して、AWS Cloud9 bash シェルで実行します。スクリプトが終了したら、以下の curl コマンド (AIAA DNS 名に /v1/models を追加) を使って、モデルが AIAA サーバーにロードされたことを確認します。

curl http://clara-LoadB-ABLH1BWGGPQ8-799966118.us-west-2.elb.amazonaws.com/v1/models

API が、先ほどロードしたモデルの概要を返します。

[{
"name": "clara_ct_seg_spleen_amp",
"labels": ["spleen"],
"description": "A pre-trained model for volumetric (3D) segmentation of the spleen from CT image",
"version": "2",
"type": "segmentation"
}]

Slicer のセットアップと自動セグメンテーションの実行

Slicer アプリケーションをインストールして起動します。拡張機能マネージャーメニューから、NvidiaAIAssistedAnnotation 拡張機能をインストールします。左上にある [Add Data] コマンドを使って、先ほど抽出した Task09_Spleen.tar アーカイブの imagesTr サブディレクトリから画像ファイルをインポートします。

[Application Settings] メニューから NVIDIA プラグインを選択し、サーバーアドレスを AIAA DNS 名と一致するように設定します。これは、上記のクイックスタート AWS CloudFormation 出力 ClaraServiceUrl から取得できます。

Slicer の [Modules] ドロップダウンから [Segment Editor] を選択し、次に [Effects] から [Nvidia AIAA] を選択します。[Auto-segmentation] ドロップダウンに、上記の設定でダウンロードした clara_ct_seg_spleen_amp モデルが表示されます。[Start] ボタンをクリックして、自動セグメンテーションジョブを開始します。

CT 画像のスライスはすべて、ロードバランサーを通じて推論のためのバックエンド Amazon ECS コンテナインスタンスに送信されます。次に、セグメンテーション結果が Slicer に返され、脾臓容積が青色で表示されます。

Clara モデルトレーニングツールの使用

NVIDIA Clara Train アプリケーションフレームワークには、コマンドラインによる強力なモデル開発ツールのセットも含まれています。セキュリティを強化するために、クイックスタートは Amazon ECS コンテナインスタンスをプライベートサブネットにデプロイします。Clara Train コンテナには、クイックスタートがデプロイした踏み台ホストへの SSH 接続を通じてアクセスします。Amazon EC2 コンソールを使って、踏み台ホストのパブリック IP と Amazon ECS コンテナインスタンスのプライベート IP を取得します。次に、SSH 認証エージェント、踏み台ホストのパブリック IP への SSH を開始して、そこから Amazon ECS コンテナインスタンスのプライベート IP への SSH を開始します。

ssh-add -K Clara2.pem
ssh -A ec2-user@35.155.36.51
ssh -A ec2-user@10.180.31.225

Amazon ECS コンテナインスタンスから docker ps を実行して、Clara Train コンテナの ID を取得します。次に、そのコンテナで対話型の bash シェルを実行します。

docker exec -it b50d10477dfc /bin/bash

コンテナのコマンドラインで、NVIDIA NCG レジストリから事前訓練済みの脾臓セグメンテーションモデルをダウンロードします。

MODEL_NAME=clara_ct_seg_spleen_amp
VERSION=1
ngc registry model download-version nvidia/med/$MODEL_NAME:$VERSION --dest /workspace/

これは、AIAA サーバーでのセグメンテーションタスクに使用したものと同じ事前訓練済みモデルですが、今回はモデルトレーニングジョブに使用されます。setup スクリプトが data ディレクトリにダウンロードしたトレーニングデータセットを抽出します。

cd /workspace/data/ 
 tar -xf ./Task09_Spleen.tar
 gunzip ./Task09_Spleen/imagesTr/*.nii.gz
 gunzip ./Task09_Spleen/imagesTs/*.nii.gz
 gunzip ./Task09_Spleen/labelsTr/*.nii.gz

AWS Cloud9 IDE 環境に戻ります。IDE 環境のディレクトリから、マウントされた Amazon EFS マウントポイントを使いやすくするためのシンボリックリンクを作成し、先ほどダウンロードしたモデルのディレクトリの所有者を変更します。

ln -s /mnt/efs/ ./workspace
sudo chown -R ec2-user ./workspace/clara_ct_seg_spleen_amp

AWS Cloud9 エディタを使って、モデルに付属する environment.json ファイルを開きます (これは ./workspace/clara_ct_seg_spleen_amp_v1/config/environment.json にあります)。DATA_ROOT を、抽出したトレーニングデータのディレクトリをポイントするように設定します。model config ディレクトリにある dataset_0.json ファイルをポイントするように DATASET_JSON を更新します。

これで、train.sh スクリプトを使ってセグメンテーションモデルのトレーニングジョブを開始する準備が整いました。

cd /workspace/clara_ct_seg_spleen_amp_v1/
chmod +x commands/train.sh
commands/train.sh

トレーニングジョブが完了したら、モデルのパフォーマンスレベルを示すモデルパフォーマンスメトリクスが報告されます。

まとめ

このブログ記事では、AWS で Clara Train アプリケーション開発フレームワークをデプロイするためのエンドツーエンドプロセスについて説明しました。また、CT スキャンのセマンティックセグメンテーションを実行するために事前訓練済みモデルを使用する方法、および独自にセグメンテーションモデルを訓練する方法のデモも行いました。この記事に基づいて、独自の AI モデルの開発、Clara Train アプリケーション開発フレームワークや、Amazon SageMaker、TensorFlow、PyTorch、または Apache MXNet での AWS のサービスの活用を開始していただければ幸いです。ヘルスケア関連のお客様で、詳細情報をご希望の方は、いつでもご遠慮なく AWS アカウントチームまでお問い合わせください。

Amazon Web Services ブログ