CNNの可視化(1)PyTorch-LightningでGrad-CAMを試す

こんばんは。座禅いぬです。最近は生成AIの話題、そしてデータサイエンスの話題で盛り上がっている一方で、ディープラーニング自体の話題はちょっと減ってきたように思います。当たり前の技術として受け入れられている結果なんでしょうね。一方で、**いくら生成AIがコードを書いてくれても実装する大変さはなくならない**ように思います。

ここ数か月は画像分類モデルをどう使うかに取り組みつついろいろな勉強を並行していますが、実際やってみると傾向として**コーディング以外の作業が多くて閉口**します。そこで重要なのが可視化であると感じています。問題や機能、効果をわかりやすくすることで、自分や他者の意思決定がスムーズになります。

NNによる画像分類で有効性が高い可視化といえば何といっても**Grad-CAM**だと個人的には思っています。素人受けもめちゃくちゃ良くて、「すごい！よくわかった！やってみよう！」につながりやすいのではないかと思います。また、僕のような初学者には前処理や画像準備の段階での間違いに気付くきっかけを与えてくれます。

**自分は初学者なので勉強しながらのメモ**という形ですが、普段つかっているPyTorch-LightningでGrad-CAMを実装する流れをまとめてみました。


## 先行記事の状況
まずはQiita上で「Grad-CAM」タグを検索してみると、たとえば以下のような記事が見つかります。

- **[Grad-CAMをPythonで実装し、予測の根拠を視覚化する](https://qiita.com/senbe/items/ee96acb8d3a49d6a98d2)**
PyTorchを前提に、標準的なResNet/VGGなどを対象としてGrad-CAMの仕組み、特徴マップと勾配の取得方法、ヒートマップ生成までを丁寧に解説されています。

- **[Grad-CAMにトライ&エラーした話](https://qiita.com/kuwagataaa/items/81f24f1342f212a6e2c7?utm_source=chatgpt.com)**
Keras／PyTorchでの動かし方やつまづきポイントを振り返る体験談。どの段階でエラーが出たか、解決策などがコードとともにまとめられています。

- **[人工知能は、アブラムシのどこを見てるんだろう](人工知能は、アブラムシのどこを見てるんだろう)**
昆虫学の専門家が、アブラムシの齢期分類を機械学習で行った際にモデルがどこに着目しているかを可視化し、ご自身の専門的知識と照らし合わせています。


これらを踏まえ今回は「**Torch-LightningでGrad-CAMを使う最小限の流れと、実装時につまづきやすいポイント**」に絞って補足します。

## Lightningでハマるポイント3選


| 罠                                                                    | 回避策                                                                     |
| -------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| `training_step / validation_step` 内は自動で `torch.no_grad()` が有効 → 勾配ゼロ | **可視化は別メソッド** (`generate_gradcam`) で実行し、`with torch.enable_grad():` で囲む |
| フックを貼る層が見つからない                                                       | ResNetなら `layer4[-1].conv2` など “最後のConv” を**1回**だけ登録                    |
| 勾配／特徴マップを前回の呼び出しが上書き                                                 | メソッド冒頭で `self.feature_maps = self.gradients = None` でリセット               |

## どういうことだってばよ
### 1. 対象モデル
- LightningModule で定義したCNNモデル（例：ResNet50など）。
- モデルは Lightning の self.model や self.backbone として実装していることを想定します。
### 2. Grad-CAMの基本原理
- 最後の畳み込み層（Conv 層）の「特徴マップ（Feature Map）」と、その層への「勾配（Gradient）」を取得。
- Grad-CAM 本来の計算手順
  1. 入力画像をモデルに通し、目的とするクラススコアを取得。
  2. そのクラススコアを微分して、最後の畳み込み層の出力に対する勾配を取得（loss.backward()） 。
  3. 取得した勾配をチャネルごとに平均化して重み（α）を計算。
  4. 特徴マップにその重みを掛けてチャネル方向に総和し、ReLU を通す。
  5. 得られたマップを元画像サイズにアップサンプリングし、ヒートマップとして可視化。

### Lightning での注意点
- LightningModule の training_step/validation_step とは別に、Grad-CAM専用の処理を行う必要がある
  通常の学習／検証ループ中に Grad-CAM を走らせると勾配フックやテンソル操作が干渉する可能性が高いため、別メソッドとして管理するのが安全です。
- 勾配（Gradient）は torch.no_grad() の文脈外で計算する
  Lightning の validation_step や test_step ではデフォルトで torch.no_grad() が有効になる場合があるため、Grad-CAM実行時には明示的に torch.enable_grad() を使って勾配を取得できる状態に切り替えます。
- フック（Hook）の登録・解除を適切に行う
  「最後の畳み込み層」を参照して、その出力（特徴マップ）と、その層への勾配（入力勾配ではなく、出力勾配）をフックでキャッチします。
  LightningModule の初期化時点でフックを登録しておけば、あとはGrad-CAMを実行したいタイミングでフック先の情報を参照できます。

## 最小限の実装例
<details><summary>サンプル（70行）を見る</summary>

```python: sample1.py
import torch, pytorch_lightning as pl
from torchvision import models
from torch import nn
import torch.nn.functional as F

class LitGradCAM(pl.LightningModule):
    def __init__(self, num_classes=2):
        super().__init__()
        self.net = models.resnet18(pretrained=True)
        self.net.fc = nn.Linear(self.net.fc.in_features, num_classes)
        self.target_layer = self.net.layer4[-1].conv2        # ←★最後のConv
        self.fmap = self.grad = None
        self._hook()

    def _hook(self):
        self.target_layer.register_forward_hook(
            lambda m, i, o: setattr(self, "fmap", o.detach()))
        self.target_layer.register_full_backward_hook(
            lambda m, gi, go: setattr(self, "grad", go[0].detach()))

    # ---- Lightning usual stuff (省略可) ----
    def forward(self, x): return self.net(x)
    def training_step(self, b, _): x,y=b; o=self(x); l=F.cross_entropy(o,y); return l
    def configure_optimizers(self): return torch.optim.Adam(self.parameters(), 1e-4)
    # ---------------------------------------

    def gradcam(self, x, cls=None):
        """x : [1,C,H,W] tensor -> 0-1 heatmap ndarray"""
        device = next(self.parameters()).device
        x = x.to(device).requires_grad_()
        self.net.eval()

        with torch.enable_grad():            # ←★no_grad() を一時解除
            out = self.net(x)
            cls = out.argmax(1).item() if cls is None else cls
            score = out[0, cls]
            self.net.zero_grad()
            score.backward(retain_graph=True)

        # αk = mean(grad_k)
        alpha = self.grad.mean(dim=(2,3), keepdim=True)
        cam = (alpha * self.fmap).sum(1, keepdim=True).relu()
        cam = F.interpolate(cam, size=x.shape[2:], mode='bilinear', align_corners=False)
        cam = cam.squeeze().cpu().numpy()
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-6)
        return cam  # (H,W)



```

</details>


## 使い方
<details><summary>サンプル(3行)を見る</summary>

```python: sample2.py
model = LitGradCAM.load_from_checkpoint("model.ckpt")  # or freshly trained
img   = preprocess("img.jpg").unsqueeze(0)             # [1,C,H,W]
heat  = model.gradcam(img)                             # ndarray

```
</details>


## まとめ：Lightning×Grad-CAMは 「分けて考える」 
1. **可視化は専用メソッドで - 学習ループから切り離す**
2. **torch.enable_grad() で勾配を取り戻す**
3. **フックは1回登録＆毎回バッファ初期化**

この3点に留意して、「モデルがどこを見ているか」→説得力とデバッグ力が跳ね上がるツールであるGrad-CAMをどんどん取り入れていきたいなと思います。

また次回以降実際に機械学習とGrad-CAMを使ってみた例について触れてみたいなと思っています。