AI・機械学習

Grad-CAMをPyTorchで実装!画像分類AIの判定根拠を可視化する方法

この記事で作るもの

以前の記事では、Grad-CAMの仕組みを解説しました(→ 「Grad-CAMとは?仕組みをわかりやすく解説」)。本記事では、その仕組みを実際にPyTorchで実装し、胸部X線画像を使って動かしていきます。

前提知識として、Grad-CAMが「勾配を使って対象クラスの判定根拠を可視化する手法である」ことを理解済みであるものとして進めます。仕組みの解説を先に読みたい方は、上記の前回記事をご覧ください。

最終的に、以下のような「元画像」と「Grad-CAMヒートマップ」を並べた出力を作ることがゴールです。

使用するライブラリは、Grad-CAMの計算部分を担うpytorch-grad-camと、モデル・画像処理を担うtransformers(Hugging Face)です。

環境準備

本記事のコードはGoogle Colaboratory上での実行を前提としています。Colab自体の基本的な使い方については、以下の記事で解説していますので、Colabに不慣れな方は先にご参照ください。

→ 「(5分でできる機械学習!Google ColabでPython初心者でもAIを動かしてみよう)」

まず、Colabの新しいノートブックを開き、以下のセルを実行して必要なライブラリをインストールします。

!pip install grad-cam transformers

次に、GPUを使用する設定にしておきます。Colabの画面上部メニューから「ランタイム」→「ランタイムのタイプを変更」を選び、ハードウェアアクセラレータを「T4 GPU」に設定してください。設定後、以下のコードでGPUが認識されているか確認します。

import torch

# GPUが使用可能かどうかを確認します

# Colabで正しくGPU設定ができていれば、「使用するデバイス: cuda」と表示されます

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(f"使用するデバイス: {device}")

モデルとデータの準備

使用モデル

本記事では、Hugging Face上で公開されているAunsiels/resnet-pneumonia-detectionを使用します。このモデルは、microsoft/resnet-50をベースに、胸部X線画像をNORMAL(正常)・PNEUMONIA(肺炎)の2クラスに分類するようファインチューニングされたものです。ライセンスはApache 2.0です。

学習に使用されたデータは、Kaggleで公開されている”Chest X-Ray Images (Pneumonia)”データセット(Kermanyらによるもの、CC BY 4.0ライセンス)です。モデルの性能は正解率0.833、F1スコア(重み付き)0.835と報告されています。

重要な注意:本記事で使用するモデルは教育目的で公開されているものであり、臨床グレードの精度検証を経たものではありません。実際の診断根拠として使用することはできません。

サンプル画像の準備

本記事で使用するサンプル画像は、以下のGitHubリポジトリに配置しています。ライセンス表記も同フォルダ内のREADMEに記載していますので、あわせてご確認ください。

本記事で使用しているサンプル画像は、匿名化済みオープンデータセット(CC BY 4.0)を使用しています。

出典: Kermany DS, Zhang K, Goldbaum M (2018), “Labeled Optical Coherence Tomography (OCT) and Chest X-Ray Images for Classification”, Mendeley Data, V2, doi:10.17632/rscbjbr9sj.2

再配布元: https://huggingface.co/datasets/hf-vision/chest-xray-pneumonia (CC BY 4.0)

本記事ではリサイズ・一部抜粋・grad-camによる加工をして再配布しています。

import urllib.request

# サンプル画像をGitHubリポジトリから取得します

image_url = " https://github.com/curiositycreates/curiosity-creates-codes/blob/main/AI/XAI/datasets/chest_xray_pneumonia/xray_fn.png?raw=true "

local_img_path = "sample_chest_xray.jpg"

urllib.request.urlretrieve(image_url, local_img_path)

print("サンプル画像のダウンロードが完了しました。")

ダウンロードされた画像ファイルは、メニューから確認することが出来ます。左のメニューバーのフォルダマークをクリックすると、ダウンロードされた「sample_chest_xray.jpg」が表示されています。

モデル構造の確認:target_layerの特定

Grad-CAMを計算するには、「モデルのどの層を対象にするか」(target_layer)を指定する必要があります。一般的には、最終畳み込み層を対象にすることが多いのですが、この「最終畳み込み層」がモデル内部のどこにあるかは、モデルの実装方法によって異なります。

たとえば、torchvisionが提供するResNet50であればmodel.layer4[-1]のように直接アクセスできますが、今回使用するモデルはHugging Faceのtransformersライブラリ経由で提供されているため、内部の階層構造が異なります

自分が使うモデルの構造が分からない場合は、named_modules()を使って実際に確認するのが確実です。

from transformers import AutoImageProcessor, AutoModelForImageClassification

model_name = "Aunsiels/resnet-pneumonia-detection"

processor = AutoImageProcessor.from_pretrained(model_name)

model = AutoModelForImageClassification.from_pretrained(model_name)

# モデル内部の層構造を確認します

# 出力の中から、最終畳み込み層に相当する部分を探します

for name, module in model.named_modules():

    print(name)

このコードを実行すると、resnet.encoder.stages以下に複数のstageとlayerが階層的に並んでいることが確認できます。今回のモデルでは、最終畳み込み層はmodel.resnet.encoder.stages[-1].layers[-1]にあたります。

# 確認した結果に基づき、target_layerを指定します

target_layers = [model.resnet.encoder.stages[-1].layers[-1]]

Hugging Faceモデルのラッパー作成

pytorch-grad-camライブラリは、モデルの出力が単純なテンソル(数値の配列)であることを前提に作られています。しかし、Hugging Faceのモデルは、通常logits(判定スコア)などを含んだ辞書型のオブジェクトを返します。このままではGrad-CAMライブラリがうまく処理できないため、「モデルを呼び出して、logitsだけを取り出して返す」ラッパークラスを用意します。

import torch.nn as nn

class HuggingFaceModelWrapper(nn.Module):

    """

    Hugging Faceモデルの出力(辞書型)から、

    Grad-CAMライブラリが必要とするlogits(テンソル)だけを

    取り出して返すためのラッパークラス

    """

    def __init__(self, model):

        super().__init__()

        self.model = model

    def forward(self, x):

        # Grad-CAM内部からの呼び出しに対して、

        # 辞書型ではなくテンソルを直接返すようにする

        outputs = self.model(pixel_values=x)

        return outputs.logits

# モデルをGPUに転送し、推論モードに設定

model = model.to(device)

model.eval()

# ラッパーでモデルを包む

wrapped_model = HuggingFaceModelWrapper(model)

Grad-CAM本体の実装

準備が整ったので、実際に画像を読み込んで推論を行い、Grad-CAMを計算していきます。

import numpy as np

from PIL import Image

from pytorch_grad_cam import GradCAM

from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget

from pytorch_grad_cam.utils.image import show_cam_on_image

import matplotlib.pyplot as plt

# ------------------------------------------

# 1. 画像の読み込み

# ------------------------------------------

image = Image.open(local_img_path).convert("RGB")

# ------------------------------------------

# 2. モデルによる推論

# ------------------------------------------

inputs = processor(images=image, return_tensors="pt")

inputs = {k: v.to(device) for k, v in inputs.items()}

with torch.no_grad():

    outputs = model(**inputs)

    logits = outputs.logits

# ソフトマックスで確率に変換し、予測クラスを確認

probs = torch.nn.functional.softmax(logits, dim=-1)

predicted_label_idx = torch.argmax(probs, dim=-1).item()

labels = model.config.id2label

predicted_label_name = labels[predicted_label_idx]

confidence = probs[0][predicted_label_idx].item()

print(f"AIの予測クラス: {predicted_label_name}")

print(f"確信度: {confidence:.2%}")

# ------------------------------------------

# 3. Grad-CAMの計算

# ------------------------------------------

# 先ほど確認したtarget_layerを対象にGrad-CAMを構築

cam = GradCAM(model=wrapped_model, target_layers=target_layers)

# 「予測されたクラス」を対象に、そのクラスの判定根拠を可視化する

targets = [ClassifierOutputTarget(predicted_label_idx)]

# Grad-CAMを実行し、ヒートマップ(0〜1の範囲の重要度マップ)を取得

cam_output = cam(input_tensor=inputs["pixel_values"], targets=targets)

grayscale_cam = cam_output[0, :]

# ------------------------------------------

# 4. 元画像とヒートマップの重ね合わせ

# ------------------------------------------

# 元画像をモデルの入力サイズにリサイズし、0〜1の範囲に正規化

# 224を決め打ちにせず、processorの設定から入力サイズを取得します

# (別のモデルに差し替えた場合でも入力サイズのズレが起きないようにするためです)

input_size = processor.size.get("shortest_edge", processor.size.get("height", 224))

input_image_np = np.array(image.resize((input_size, input_size))) / 255.0

# ヒートマップを元画像に重ね合わせる

visualization = show_cam_on_image(input_image_np, grayscale_cam, use_rgb=True)

# ------------------------------------------

# 5. 結果の表示

# ------------------------------------------

plt.figure(figsize=(12, 6))

plt.subplot(1, 2, 1)

plt.imshow(image)

plt.title("Original Image")

plt.axis("off")

plt.subplot(1, 2, 2)

plt.imshow(visualization)

plt.title(f"Grad-CAM (Target: {predicted_label_name})")

plt.axis("off")

plt.tight_layout()

plt.show()

出力を確認する

上記のコードを実行すると、元画像と、Grad-CAMのヒートマップを重ね合わせた画像が並んで表示されます。

肺炎陽性の画像であれば、肺野の陰影がある領域に注目が集まっていることが確認できるはずです。この結果の解釈(正しく判定できた場合・誤判定した場合の違いや、そこから見えてくる限界)については、前回の記事で誤判定例も含めて扱っていますので、あわせてご覧ください。

「Grad-CAMとは?仕組みをわかりやすく解説」(誤判定事例・限界について解説)

まとめ

本記事では、Hugging Face上の肺炎判定モデルを題材に、PyTorchでGrad-CAMを実装する手順を解説しました。ポイントは以下の3点です。

  • モデルの内部構造はnamed_modules()で確認し、target_layerを特定する
  • Hugging Faceモデルの辞書型出力は、ラッパークラスでテンソルに変換する必要がある
  • 実装したコードはColab上でそのまま実行可能

今後は、Grad-CAMの発展形であるGrad-CAM++・Score-CAM・Layer-CAMについても解説記事を作成予定です。

本記事で使用したコードとサンプル画像は、以下のGitHubリポジトリで公開しています。

→ GitHubリポジトリ:curiosity-creates-codes(作成中)


本記事で使用しているモデル・データは教育目的で公開されているものであり、実際の臨床診断の根拠として使用できる精度検証を経たものではありません。

GRAD-CAMの実装と関連がある記事

Python初心者はいきなりインストール不要|Google Colabを勧める理由

5分でできる機械学習!Google ColabでPython初心者でもAIを動かしてみよう

Grad-CAMとは?仕組みをわかりやすく解説|胸部X線AIで見る判断根拠の可視化と限界

コメント

タイトルとURLをコピーしました