Amazon SageMaker で高解像度胸部 X 線医療画像を分類する

医療画像処理は、深層学習が大きな効果を発揮する重要な分野の 1 つです。一般的な処理としては、様々な医用画像モダリティを用いた分類、検出、セグメンテーションがあります。このブログ記事では、HIPAA 準拠サービスである Amazon SageMaker を使用して、Amazon SageMaker 画像分類アルゴリズムで胸部 X 線画像分類の深層学習モデルをトレーニングする方法を概説します。この画像分類アルゴリズムが、高解像度医療画像を分析するための有効なツールとなり得ることを示したいと考えています。マルチラベルサポートや混合精度トレーニングなどのアルゴリズムの新しい機能を使用して、 混合精度モードを使用する胸部 X 線画像分類モデルが Amazon EC2 P3 インスタンスで float32 データタイプを使用する場合より、トレーニングによって 33% 速くなることを示します。また、高解像度の画像で胸部 X 線画像をどのようにトレーニングすることができるかを示し、低解像度のモデルでもパフォーマンスを向上できることを示します。

高解像度胸部 X 線医療画像の分類

ディープニューラルネットワークベースのアプローチは、通常は、メモリの制約のために低解像度の画像で動作します。画像分類 (ResNet-152) で使用する一般的なディープネットワークでは、256×256 サイズの画像でも大きなメモリが必要です。また、メモリ要件は、トレーニングで使用するバッチサイズにも依存します。ただし、一部の疾患は、胸部 X 線画像の小さな領域にしか存在しないため、高解像度画像分類の恩恵を受ける可能性が高くなります。

胸部 X 線画像のデータセットは、アメリカ国立衛生研究所 (NIH) [1] から公表されており、https://nihcc.app.box.com/v/ChestXray-NIHCC から入手できます。このデータセットは、30,805 人の患者からの 112,120 件の正面図 X 線画像で構成されます。これらの画像は、自然言語処理を使用して、関連付けられた放射線レポートからマイニングされた最大 14 のテキストマイニング病気画像ラベルを含むことができます。これらの 14 のラベルのいずれも、それぞれの X 線に関連付けることができることが、マルチラベル画像分類の問題となります。疾患のラベルは、無気肺、浸潤影、浸潤、気胸、浮腫、気腫、線維症、滲出、肺炎、胸膜肥厚、小瘤、結節、腫瘤およびヘルニアです。

近年、胸部 X 線画像分類に深層学習アルゴリズムが適用されています [1]。ImageNet データセットで事前にトレーニングした 50 層の ResNet を使用して、胸部 X 線画像による疾患分類をトレーニングしました。画像は 224×224 rの解像度にサイズ変更され、マルチラベル設定で加重クロスエントロピー損失を使用してトレーニングされました。最近では、448×448 の解像度でトレーニングを行ったときに、特に腫瘤や結節などのクラスでパフォーマンスの改善が示されています [2]。

より高い解像度で画像分類を行うには、いくつかの方法があります。バッチサイズを減らしてメモリ要件を満たしたり、GPU/ マシン間でトレーニングを分散したりすることができます。バッチサイズが小さいとノイズの多いグラデーションがローカルコンバージェンスにつながるため、バッチサイズを小さくすると精度が低下する可能性があります。類似の設定を使用して、Amazon SageMaker 画像分類アルゴリズムの新機能を使用して最初に結果を再現し、フル解像度のトレーニングがパフォーマンスを向上させる方法を示します。Amazon SageMaker 画像分類アルゴリズムを強化することで、複数のマシンを使用してバッチサイズを減らすことなく、マルチラベル設定で高解像度画像をより迅速にトレーニングする方法を紹介します。

Amazon SageMaker による画像の分類

Amazon SageMaker 画像分類は、Amazon SageMaker の組み込みアルゴリズムの 1 つです。ResNet アーキテクチャを実装し、フルトレーニングと転移学習の両方をサポートしています。転移学習に使用される事前トレーニング済みのモデルは、ImageNet データセットを使用してトレーニングされています。このアルゴリズムは、異なるデータセットに対応する ResNet アーキテクチャの複数の深度をサポートします。転移学習モードでは、アルゴリズムは最終的に完全に接続されたレイヤーを、指定された出力数を持つ新しい完全接続のレイヤーに置き換え、ウェイトを初期化し、ネットワーク全体を新しいデータセットで微調整します。このアルゴリズムは、シングルおよびマルチ GPU セットアップとシングルまたはマルチマシントレーニングをサポートします。マルチマシントレーニングでは、アルゴリズムはウェイトの同期および非同期の両方の更新をサポートします。

最近、Amazon SageMaker 画像分類アルゴリズムにいくつかの新機能を追加しました。まず、ネットワークの計算が float16 モードで行われ、同様の精度を維持しながらメモリを少なくしてトレーニングを行う、混合精度モードでのトレーニングを可能にしました。また、各画像を複数のクラスに分類できるマルチラベルデータセットのトレーニング用のマルチラベル入力サポートも有効にしました。さらに、クラスの不均衡に対処するために、加重損失更新の使用を有効にします。

マルチラベル画像分類のためのデータの準備

それぞれの胸部 X 線画像のオリジナルは、1024×1024 のサイズです。最初の結果を得るために 224×224 の入力解像度を使用し、その後、完全な解像度で結果を得ました。入力データセットは、これらのセット間で患者が重複しないトレーニング/検証およびテストに分割され、14 の疾患はマルチラベル分類の 14 のバイナリ出力としてモデル化されます。それぞれの疾患出力は疾患の有無を示し、「疾患なし」はすべてのラベルがゼロとして表現されます。最初に、患者 ID に基づいてトレーニング、検証、テストの 3 つのセットにデータを分割しました。

import panda as pd
import numpy as np

trainper = 0.7
valper = 0.1
file_name = 'Data_Entry_2017.csv'

a = pd.read_csv(file_name)
patient_ids = a['Patient ID']
uniq_pids = np.unique(patient_ids)
np.random.shuffle(uniq_pids)
total_ids = len(uniq_pids)

trainset = int(trainper*total_ids)
valset = trainset+int(valper*total_ids)
testset = trainset+valset

train = uniq_pids[:trainset]
val = uniq_pids[trainset+1:valset]
test = uniq_pids[valset+1:]
print('Number of patient ids: training: %d, validation: %d, testing: %d'%(len(train), len(val), len(test)))

traindata = a.loc[a['Patient ID'].isin(train)]
valdata = a.loc[a['Patient ID'].isin(val)]
testdata = a.loc[a['Patient ID'].isin(test)]

traindata.to_csv('traindata.csv', sep=',', header=False, index=False)
valdata.to_csv('valdata.csv', sep=',', header=False, index=False)
testdata.to_csv('testdata.csv', sep=',', header=False, index=False)

次に、各疾患がバイナリラベルのセットにマッピングされるマルチラベル分類用の「lst」ファイルを作成します。

import csv

def gen_set(csvfile, outputfile):
    disease_list = ['Atelectasis', 'Consolidation', 'Infiltration', 'Pneumothorax', 'Edema', 'Emphysema', \
                   'Fibrosis', 'Effusion', 'Pneumonia', 'Pleural_Thickening', 'Cardiomegaly', 'Nodule', 'Mass', \
                   'Hernia']
    alldiseases = {disease:i for i,disease in enumerate(disease_list)}
    with open(outputfile, 'w') as fp:
        with open(csvfile, 'r') as cfile:
            line = csv.reader(cfile, delimiter=',')
            index = 0
            for element in line:
                # 最初の列は画像ファイル名、
                # 2 番目の列は「|」で区切られた疾病の名前のリスト
                diseases = element[1].split('|')
                fp.write('%d\t'%index)
                for d in alldiseases:
                    if d in diseases:
                        fp.write('%d\t'%1)
                    else:
                        fp.write('%d\t'%0)
                fp.write('images/%s\n' % element[0])
                index += 1
                 
gen_set('traindata.csv', 'chestxraytrain.lst')
gen_set('valdata.csv', 'chestxrayval.lst')
gen_set('testdata.csv', 'chestxraytest.lst')      

次に、トレーニングと検証のために im2rec.py を使用して、recordio ファイルを作成します。recordio ファイルがマルチラベルファイルとして作成されるように、pack_label オプションを渡します。詳細については、Amazon SageMaker マルチラベル画像分類ノートを参照してください。

python im2rec.py --pack-label chestxraytrain.lst .
python im2rec.py --pack-label chestxrayval.lst .

胸部 X 線データセットの画像分類結果

ResNet-50 モデルを使用し、最初に 224×224 の入力イメージサイズでネットワークをトレーニングしました。ランダムクロッピングや画像変換などのデータ拡張手法を使用しました。胸部 X 線画像が ImageNet 画像と異なる場合でも、ImageNet データセットで事前にトレーニングされたモデルを使用すると、分類精度が向上します。そのため、Amazon SageMaker 画像分類アルゴリズムの use_pretrained_model ハイパーパラメータを使用してネットワークをトレーニングしました。これはマルチラベル分類であるため、multi_label パラメータを 1 に設定しました。ネットワークが 224×224 の領域を入力画像からトリミンするグことができるように、胸部 X 線画像をトレーニング前に 256 にサイズ変更しました。

次のコードスニペットは、Amazon SageMaker Estimator インターフェイスと画像分類アルゴリズムを使用して、これをどのように行うことができるかを示しています。

import sagemaker
from sagemaker import get_execution_role
from sagemaker.amazon.amazon_estimator import get_image_uri

role = get_execution_role()
sess = sagemaker.Session()
bucket=sess.default_bucket()
prefix = 'ic-chestxray'

training_image = get_image_uri(sess.boto_region_name, 'image-classification', repo_version="latest")
s3train = 's3://{}/train/'.format(bucket)
s3validation = 's3://{}/validation/'.format(bucket)
s3_output_location = 's3://{}/{}/output'.format(bucket, prefix)
multilabel_ic = sagemaker.estimator.Estimator(training_image, role,
                        train_instance_count=1,
                        train_instance_type='ml.p3.16xlarge',
                        train_volume_size = 50, train_max_run = 360000,
                        input_mode= 'File', output_path=s3_output_location,
                        sagemaker_session=sess)
multilabel_ic.set_hyperparameters(num_layers=50, use_pretrained_model=1,
                                        image_shape = "3,224,224", num_classes=14,
                                        mini_batch_size=256,
                                        resize=256,  epochs=100,
                                        learning_rate=0.0005, optimizer='adam',
                                        num_training_samples=80000,
                                        augmentation_type = 'crop_color_transform',
                                        precision_dtype='float32', multi_label = 1)
train_data = sagemaker.session.s3_input(s3train, distribution='FullyReplicated',
                                                content_type='application/x-recordio',
                                                s3_data_type='S3Prefix')
validation_data = sagemaker.session.s3_input(s3validation, distribution='FullyReplicated',
                                                content_type='application/x-recordio',
                                                s3_data_type='S3Prefix')
data_channels = {'train': train_data, 'validation': validation_data}
multilabel_ic.fit(inputs=data_channels, logs=True)

ネットワークの性能をテストセットでテストし、平均 AUC と共に個々の AUC を計算しました。結果は次の表にまとめられています。得られた平均 AUC は 0.768 でした。

加重損失を伴うトレーニング

画像分類に導入された追加機能としては、クラスの不均衡に対処するための加重損失の使用があります。通常、マルチラベルのデータセットでトレーニングする場合、クラス間に不均衡が存在する可能性があります。この不均衡によって、ネットワーク学習がある別のクラスよりも特定のクラスを学習することにつながることがあります。こうした事態を避けるために、Amazon SageMaker 画像分類アルゴリズムは use_weighted_loss ハイパーパラメータを使用してサンプルのバランスをとっています。このパラメータを 1 に設定すると、トレーニングセット内のそのラベルのサンプル数に基づいて、各ラベルの重み値が計算されます。まず、各クラスのサンプル数がトレーニングセットから計算され、損失更新の重みがそのクラスについて N/N_l に設定されます。ここで、N はトレーニングセット内のサンプルの総数であり、N_l はトレーニングセット内のクラス l のサンプルの総数です。これにより、それぞれのクラスに対して重みに基づいて勾配の更新を異なる重み付けで計算し、バランスのとれたトレーニングを可能にします。平均 AUC は、引き続き 224×224 の入力解像度を使用しながら、重み付き損失機能を有効にしてトレーニングした場合、0.814 に増加しました。

multilabel_ic.set_hyperparameters(num_layers=50, use_pretrained_model=1,
                                        image_shape = "3,224,224", num_classes=14,
                                        mini_batch_size=256, resize=256,  epochs=100,
                                        learning_rate=0.0005, optimizer='adam',
                                        num_training_samples=80000, use_weighted_loss=1,
                                        augmentation_type = 'crop_color_transform',
                                        precision_dtype='float32', multi_label = 1)

混合精度でのトレーニング

Amazon SageMaker 画像分類アルゴリズムは、現在、混合精度モードでのトレーニングをサポートしています。これは、「float32」 (デフォルト) または「float16」に設定できるハイパーパラメータ precision_dtype によってコントロールされます。混合精度モードでは、ネットワークは、高精度 (float32) でマスターウェイトを維持しながら、低精度 (float16) でバックワードおよびフォワードのパスを計算します。これにより、同様の精度を維持しながらトレーニングを高速化することができます。混合精度モードを使用することにより、トレーニング時間は 33% 短縮されますが、全体の AUC は 0.821 であり、これは float32 トレーニングで得られた値と同様です。トレーニング時間の短縮は、高解像度入力 (次のセクションを参照) に 2 つのインスタンスを使用してトレーニングを行った場合に増加し、47% に達しました。

multilabel_ic.set_hyperparameters(num_layers=50, use_pretrained_model=1,
                                        image_shape = "3,224,224", num_classes=14,
                                        mini_batch_size=256, resize=256, epochs=100,
                                        learning_rate=0.0005, optimizer='adam',
                                        num_training_samples=80000, use_weighted_loss=1,
                                        augmentation_type = 'crop_color_transform',
                                        precision_dtype=’float16’, multi_label = 1)

高解像度入力によるトレーニング

image_shape パラメータを 896×896 に設定して、元の入力解像度を使用しました。このトレーニングでは、use_weighted_loss 機能と float32 精度を使用しました。この解像度を使用した理由は、ネットワークがデータ拡張中に 1024×1024 から 896×896 の領域をサンプリングできるためです。高解像度はより多くのメモリを使用するので、一般には、batch_size を下げてネットワークをトレーニングします。ただし、Amazon SageMaker 画像分類は分散トレーニングをサポートしているので、複数のインスタンス間でトレーニングを実行することで batch_size を維持することができました。これは、Amazon SageMaker トレーニングで instance_count parameter パラメータを 2 に設定することで行われます。この解像度の平均 AUCは、特に高解像度入力の恩恵を受ける可能性がある結節などのクラスでは 0.830 に増加しました。mixed_precision を 1 に設定してトレーニングを行ったところ、平均 AUC は 0.825 でした。トレーニングは以前と同じコードを使用していましたが、train_instance_count = 2、image_shape=”3,896,896” と設定し、resize パラメータは設定しませんでした。

multilabel_ic = sagemaker.estimator.Estimator(training_image, role, train_instance_count=2,
                                                train_instance_type='ml.p3.16xlarge',
                                                train_volume_size = 50, train_max_run = 360000,
                                                input_mode= 'File', output_path=s3_output_location,
                                                sagemaker_session=sess)

multilabel_ic.set_hyperparameters(num_layers=50, use_pretrained_model=1,
                                        image_shape = "3,896,896", num_classes=14,
                                        mini_batch_size=64, epochs=100,
                                        learning_rate=0.00025, optimizer='adam',
                                        num_training_samples=80000, use_weighted_loss=1,
                                        augmentation_type = 'crop_color_transform',
                                        precision_dtype='float32', multi_label = 1)

まとめ

このブログ記事では、Amazon SageMaker 画像分類アルゴリズムの新機能を使用して、高解像度胸部 X 線データセットをトレーニングする方法を示しました。特に、multi_label パラメータを使用して複数ラベルのデータセットを学習する方法、およびラベル内のサンプルバランシングの use_weighted_loss パラメータを表示して全体の精度を向上させる方法を示しました。また、同様の精度を維持しながら、複合精度を使用してトレーニング時間を 33% 短縮する方法と、複数のインスタンスを使用して、トレーニング時間をさらに 47% 短縮する方法も示しました。さらに、このアルゴリズムを使用して高解像度画像によるトレーニングを行う方法を示し、低解像度入力を使用する場合に比べて精度が向上しました。

参考文献

1. X.Wang, Y. Peng, L. Lu, Z. Lu, M. Bagheri, R. M.Summers.ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases.CVPR 2017
2. Ivo M. Baltruschat, Hannes Nickisch, Michael Grass, Tobias Knopp, and Axel Saalbach, Comparison of Deep Learning Approaches for Multi-Label Chest X-Ray Classification, arxiv 1803.02315 (https://arxiv.org/abs/1803.02315)

著者について

Gurumurthy Swaminathan は、Amazon AI プラットフォームグループの上級応用サイエンティストであり、Sagemaker のコンピュータビジョンアルゴリズムの構築に取り組んでいます。彼の現在の研究領域には、ニューラルネットワーク圧縮とコンピュータビジョンアルゴリズムが含まれます。

Vineet Khare は、AWS Deep Learning のサイエンスマネージャーです。彼は、研究の最前線にある技術を使用して、AWS 顧客の人工知能および機械学習のアプリケーションを構築することに重点を置いています。余暇には、彼は読書、ハイキング、家族との時間を楽しんでいます。