CAM系手法のちょっとしたまとめとPyTorch実装

深層学習ベースの画像認識手法について判断根拠を与える手法のうちCAM系のものについて勉強し、モデルの実装自身を変更せずに実装でき、有名な手法を選んでPyTorchの実装練習もかねて実装してみました。

実装について
Pytorch 実装上のtips
- hook機能
class activation mapping (CAM) [Zhou+, CVPR2016]
- アイデア
- heatmapの例
Grad-CAM [Selvaraju+, ICCV2017]
- アイデア
- heatmapの例
Grad-CAM++ [Chattopadhay+, WACV2018]
- アイデア
- heatmapの例
Score-CAM [Wang+, CVPRW2020]
- アイデア
- heatmapの例
参考サイトなど

実装について

colab.research.google.com
誰でも試せるように、Google Colabで実装を行いました。
使用するデータセットは、STL-10と呼ばれる画像のデータセットです。元の画像が96x96と比較的大きく、10クラスしかないためちゃんと実装できてるかの検証が比較的しやすく、各クラス500枚しかないためColabでも短時間で学習が終わりそうというのが選んだ理由です。
モデル自体の識別性能は意識せず、全体的に自分で実装したかったのでモデルは自分で適当に作ったものです。data augmentationや正規化もしていません。そのため検証データでの性能がそれほど良くないですが、目的はheatmapをうまく出力するということなのでご容赦ください。
今後いろんなCAMでのheatmap作成結果を載せていきますが、ラベルが「車」でちゃんと「車」と識別した図1のvalidデータの画像についてheatmapを作ることにします。

f:id:verliezer93764:20210220114751p:plain — 図1: heatmap作成時の入力画像

図1の画像を識別モデルに入力した際の各クラスの出力をsoftmaxで確率分布化した結果を簡単に載せておきます。carが約0.85, truckが約0.15ですね。

      class   probability
0  airplane  1.360998e-05
1      bird  4.917460e-08
2       car  8.490953e-01
3       cat  7.586426e-07
4      deer  1.025798e-07
5       dog  3.446858e-07
6     horse  2.043831e-05
7    monkey  1.075173e-07
8      ship  1.048932e-05
9     truck  1.508588e-01

Pytorch 実装上のtips

hook機能

各層の勾配を求めるためにはtorch.autograd.grad()も使えますが、各層の特徴マップと勾配を求めるにはhook機能が便利です。
CAMでは、こんな感じの実装をしています。

class CAM():
  def __init__(self, model):
    self.model = model
    self.image_reconstruction = None
    self.layer_names = []
    self.feature_maps_each_layer = []
    self.feature_gradients_each_layer = []
    self.forward_hook_handles = []
    self.backward_hook_handles = []
    self.model.eval()
    self.register_hooks()

  def forward_hook(self, module, inp, outp):
    self.feature_maps_each_layer.append(outp)

  def backward_hook(self, module, inp, outp):
    self.feature_gradients_each_layer.append(outp)

  def register_hooks(self):
    for name, layer in self.model._modules.items():
      if isinstance(layer, nn.Conv2d):
        self.layer_names.append(name)
        self.forward_hook_handles.append(layer.register_forward_hook(self.forward_hook))
        self.backward_hook_handles.append(layer.register_backward_hook(self.backward_hook))
  
  def remove_hooks(self):
    for i in range(len(self.forward_hook_handles)):
      self.forward_hook_handles[i].remove()
    for i in range(len(self.backward_hook_handles)):
      self.backward_hook_handles[i].remove()

  def analyze(self, target_image):
    target_image_numpy = target_image.squeeze(dim=0).to('cpu').detach().numpy().transpose((1,2,0))
    probs = self.model.forward(target_image, analyze_cam=True)
    self.remove_hooks()
    df = pd.DataFrame({'class': LABELS, 'probability': F.softmax(probs, dim=1).squeeze(dim=0).to('cpu').detach().numpy()})
    print(df)
    
    # 各クラスについてheatmapを求める
    heatmaps = []
    for i in range(len(LABELS)):
      heatmap = sum([self.feature_maps_each_layer[-1][:,j] * self.model.fc.weight[i][j] for j in range(512)]).to('cpu').detach().numpy().transpose((1,2,0))
      heatmaps.append(heatmap)
    return heatmaps

CAMクラスのインスタンスが作成されると、__init__()によりself.register_hooks()が実行されます。
self.register_hooks()内では、各層がConv2d層であるかどうかを調べ、Conv2d層であった場合にはその層についてregister_forward_hook()およびregister_backward_hook()を実行する (フックをかける) というような処理をしています。そうすると、各Conv2d層について、前者はforward-passを行った場合、後者はbackward-passを行った場合に自動的に実行されるようになります。register_backward_hook()は、3つの引数module, inp, outを持ち、forward-passを行うと、inpにはその層への入力、outにはその層の処理を行った後の出力が自動的に入り、関数内の処理で使うことができます。register_backward_hook()も3つの引数module, inp, outを持ち、backward-passを行うと、inpにはモデルの出力に対するその層の勾配、outにはその層への入力に対するその層の出力の勾配が入り、関数内の処理で使うことができます。連鎖律によって、前の層のoutと次の層のinpは等しくなります。
ただ、hookをかけっぱなしだと、上の実装では層の入力などを保持するリストにforward-passのたびに記録されていきColabでメモリを圧迫してくるなどの不都合が起きます。そこで、remove_hook()によってフックを外す処理も入れておいた方がいいです。
このhook機能を使えば、各層の入力、出力、勾配を容易に獲得できます。参考サイトにものせてありますが、Understanding Pytorch hooks | Kaggleでの解説が非常にわかりやすいのでおすすめします。

class activation mapping (CAM) [Zhou+, CVPR2016]

アイデア

CAM系の手法は、heatmapを作成することで、モデルが画像のどこの部分に注目して識別を行ったのか判断根拠を提示する手法です。
CAMでは、最後のConv2d層で得られた各特徴マップについて画素の平均値をとって平坦化する global average pooling (GAP) を適用し、その出力である特長マップの個数と同じ要素数のベクトルを全結合層に入力することによって識別クラス数の要素数のベクトルを出力して予測するようなモデルに限って使用することができます。
いま、クラス $c$ のheatmapを作りたいとして、GAP前の最後の特徴マップの枚数を $n$ 、 $k (1\leq k\leq n)$ 番目の特徴マップを $A^{k}$ 、 $A^{k}$ の座標 $(i, j)$ における値を $A^{k}_{ij}$ 、全結合層の、GAP後のベクトルの特徴マップ $A^{k}$ に対応する要素から出力層のクラス $c$ のノードへの重みを $w^{c}_{k}$ と表します。

f:id:verliezer93764:20210217102309p:plain — 図2: GAPを用いた識別器の概要

CNNでは特徴マップの各画素の輝度値を注目画素の周辺画素とフィルターの畳み込みから求めていくため、特徴マップの各画素の位置と入力画像の各画素の位置は対応しているといえます。これがCAMのポイントその1です。次に、 $k$ 番目の特徴マップ $A^{k}$ の平均から出力層のクラス $c$ のノードへの重み $w^{c}_{k}$ が大きいほどクラス $c$ に対する出力が変化し識別に影響を与えることから、重み $w^{c}_{k}$ は $k$ 番目の特徴マップ $A^{k}$ の、予測に対する寄与の大きさを表しているといえます。これがCAMのポイントその2です。
これらのポイントを踏まえ、CAMのクラス $c$ に対するheatmapの座標 $(i,j)$ の値 $CAM^{c}_{ij}$ は次の式で表せます。
$\displaystyle CAM_{ij}=\sum_{k} A^{k}_{ij}w^{c}_{k} \cdots (1)$
すなわち、heatmapは特徴マップ $A^{k}$ の $w^{c}_{k}$ による重み付き和といえます。特徴マップの各画素の位置と入力画像の各画素の位置は対応しているため、heatmapの座標 $(s,t)$ の値が大きいとき、入力画像 $F$ におけるその座標付近の画素がモデルの予測に大きな影響を与えているといえます。言い換えれば、モデルはその座標付近の値を根拠に予測を行ったといえます。
CAMは特徴マップと重みを用いるだけで計算できるので、他のCAM系手法と比べてシンプルで実装が簡単です。一方で、CAMはGAPを用いたモデルに対してしか使用できず、汎用性に欠けます。

heatmapの例

各クラス $c$ で、heatmapを作ってみました。carとtruckクラスについては画像下部の部分が赤くなっており、その部分を根拠に予測結果を出したことがわかります。

f:id:verliezer93764:20210220115127p:plain — 図3: CAMによる各クラスでのheatmap出力結果

Grad-CAM [Selvaraju+, ICCV2017]

アイデア

CAMは、GAPの結果から出力層に至る重みを特徴マップの予測に対する寄与の大きさとして使用していましたが、これではGAPを利用したモデルに応用が限定されてしまいます。Grad-CAMでは、クラス $c$ のsoftmax関数適用前の出力 $Y^{c}$ からbackpropにより求められる特徴マップ上の各画素の勾配の平均を特徴マップの予測に対する寄与の大きさとして使用します。クラス $c$ の出力に対する勾配の平均が大きければ、その特徴マップが微小に変化しただけで予測結果が大きく変わる、すなわち予測結果に大きく寄与するといえます。使用する情報は特徴マップと特徴マップの各画素の勾配情報であるため、GAPを利用したモデルに限定されない任意のCNNモデルに応用でき、任意の層での可視化も可能になります (ただし最終のConv2d層が一番判断根拠を示すのでわざわざ各層でheatmapを出す必要はありません)。
これらを踏まえて、 $y^{c}$ に対する勾配の平均 $\alpha^{c}_{k}$ を求めます。特徴マップの画素数を $Z$ とします。
$\displaystyle \alpha^{c}_{k}=\frac{1}{Z} \sum_{i,j}\frac{\partial Y^{c}}{\partial A^{k}_{ij}} \cdots (2)$
この $\alpha^{c}_{k}$ を用いて、Grad-CAMのクラス $c$ に対するheatmapの座標 $(i,j)$ の値 $GradCAM^{c}_{ij}$ は次の(3)式で表せます。CAMと同じで、特徴マップの重み付き和です。ReLUを用いることで、正の勾配の画素のみを扱います。勾配が正であれば、その画素の輝度値が高いほどクラス $c$ に対する出力が大きくなり、予測結果に寄与するためです。
$\displaystyle GradCAM^{c}_{ij}=ReLU \left( \sum_{k} A^{k}_{ij} \alpha^{c}_{k} \right) \cdots (3)$
ところで、Grad-CAMはその名の通りCAMの拡張版といえます。いま、モデルが図2のような構成だったとすると、クラス $c$ のsoftmax関数適用前の出力 $Y^{c}$ は(4)式で表せます。
$\displaystyle Y^{c}=\sum_{k} w^{c}_{k} \frac{1}{Z}\sum_{i,j} A^{k}_{ij} \cdots (4)$
(4)式において、 $\displaystyle \frac{1}{Z}\sum_{i,j} A^k_{ij}$ は特徴マップのGAPであり、それの重み付き和がクラス $c$ のsoftmax関数適用前の出力 $Y^{c}$ となっています。
このGAPの部分を $F^{k}$ とします。
$\displaystyle F^{k}=\frac{1}{Z}\sum_{i,j} A^{k}_{ij} \cdots (5)$
$\displaystyle Y^{c}=\sum_{k} w^{c}_{k} F^{k} \cdots (6)$
次に、 $Y^{c}$ の $F^{k}$ による微分を考えます。
$\displaystyle \frac{\partial Y^{c}}{\partial F^{k}}=\frac{\frac{\partial Y^{c}}{\partial A^{k}_{ij}}}{\frac{\partial F^{k}}{\partial A^{k}_{ij}}} \cdots (7)$
(6)式の両辺を $F^{k}$ で微分すると $\frac{\partial Y^{c}}{\partial F^{k}}=w^{c}_{k}$ 、(5)式の両辺を $A^{k}_{ij}$ で微分すると $\frac{\partial F^{k}}{\partial A^{k}_ij}=\frac{1}{Z}$ であり、これらを(7)に代入して(8)式を得ます。
$\displaystyle w^{c}_{k}=Z\frac{\partial Y^{c}}{\partial A^{k}_{ij}} \cdots (8)$
(8)式について、すべての(i, j)の総和をとります。
$\displaystyle \sum_{i,j} w^{c}_{k}=\sum_{i,j} Z\frac{\partial Y^{c}}{\partial A^{k}_{ij}} \cdots (9)$
$\sum_{i,j}=Z$ であるから、
$\displaystyle Zw^{c}_{k}=Z\sum_{i,j} \frac{\partial Y^{c}}{\partial A^{k}_{ij}}$
すなわち
$\displaystyle w^{c}_{k}=\sum_{i,j} \frac{\partial Y^{c}}{\partial A^{k}_{ij}} \cdots (10)$
したがって、CAMにおける重みを、特徴マップの各画素の勾配の総和で代替した形となっています。平均ではなく総和となっていますが、可視化の過程で正規化を行うことになるので平均をとっても総和をとってもheatmapは変わりません。

heatmapの例

Grad-CAMでは任意の層でheatmapを作成できるため、carクラスに関して各Conv2d層でheatmapを作成してみました。最終層から前の層になるにつれて赤い部分が分散していき、最初の層付近では画像のエッジ付近のみ赤くなっていることがわかります。

f:id:verliezer93764:20210220120006p:plain — 図4: Grad-CAMによる各層でのheatmap出力結果

Grad-CAM++ [Chattopadhay+, WACV2018]

アイデア

特徴マップの各画素の勾配の平均値を特徴マップの予測結果への寄与の程度としているGrad-CAMでは、局所的な特徴をとらえた特徴マップが軽視されてしまうという問題があります。たとえば、識別クラスに含まれる物体が画像上のサイズが異なっている状態で2つ入力画像に映りこんでいたとして、大きく映っているほうをA1、小さく映っている方をA2とします。また、CNNの最後の特徴マップが2枚あり、対象クラスの出力に関するbackpropを行ったところ1枚目が物体A1、2枚目が物体A2に反応していた (その部分の勾配が大きかった) とします。このとき、Grad-CAMでは、特徴マップの各画素の勾配の平均値を特徴マップの予測結果への寄与の程度としていたことから、多くの画素で勾配の大きかった1枚目は重視され、少ない画素で勾配が大きかった2枚目は重視されないということになり、作成されるheatmapは「物体Aのみを判断根拠として予測結果を出力した」ことが読み取れるようなものとなってしまいます。
これを解決しようとしたのがGrad-CAM++です。Grad-CAM++は、特徴マップの各画素の勾配の平均値を特徴マップの予測結果への寄与の程度とせず、各画素の勾配の重み付き和を特徴マップの予測結果への寄与の程度とします。クラス $c$ に関する、この特徴マップ $k$ の画素ごとの重みを $\alpha^{kc}_{ij}$ として、特徴マップ $k$ に対する重み $w^{c}_{k}$ を次の式で求めます。
$\displaystyle w^{c}_{k}=\sum_{i,j} \alpha^{kc}_{ij} \cdot ReLU(\frac{\partial Y^{c}}{\partial A^{k}_{ij}}) \cdots (11)$
重み $\alpha^{kc}_{ij}$ は、(11)式を(4)式に代入し、二次微分をとって $\alpha^{kc}_{ij}$ について解くことによって獲得できます。(4)式から重み $\alpha^{kc}_{ij}$ を求めるということは、Grad-CAM++はGrad-CAMと同様にCAMの拡張であり、かつGrad-CAM++は最後のConv2d層にしか適用できません。
$\displaystyle \alpha^{kc}_{ij} = \frac{\frac{\partial^{2} Y^{c}}{(\partial A^{k}_{ij})^{2}}}{2 \cdot \frac{\partial^{2} Y^{c}}{(\partial A^{k}_{ij})^{2}} + \sum_{a,b} A^{k}_{ab} \frac{\partial^{3} Y^{c}}{(\partial A^{k}_{ij})^{3}}} \cdots (12)$
(12)式から、 $\alpha^{kc}_{ij}$ を求めるにはクラス $c$ のsoftmax関数適用前の出力 $Y^{c}$ の $A^{k}_{ij}$ に関する二次微分と三次微分を求める必要がありますが、PyTorchの自動微分を利用してもこれらを求めるのは容易ではありません。しかし、最後の特徴マップから出力層に至るまでにReLUなどの線形変換しか存在しない場合、高次微分は一次微分のみを用いて簡単に求めることができることが論文内で述べられています。
ややこしい話になりますが、今の $Y^{c}$ を $S^{c}$ と表し、新たに $Y^{c}=exp(S^{c})$ をsoftmax適用前の出力とします。すなわち、自然対数の底 $e$ を底とする指数関数による出力を新たな出力とします。これにより各クラスの出力の値が変化してしまいますが、exp(x)は単調増加関数であるため各クラスの出力の大小関係は変化せず、識別結果に影響しないため良しとします。
それでは $Y^{c}$ の $A^{k}_{ij}$ に関する一次微分、二次微分、三次微分を求めていきます。
まず、一次微分は合成関数の微分法により(13)式で求まります。
$\displaystyle \frac{\partial Y^{c}}{\partial A^{k}_{ij}} = exp(S^{c}) \frac{\partial S^{c}}{\partial A^{k}_{ij}} \cdots (13)$
$S^{c}$ の一次微分 $\frac{\partial S^{c}}{\partial A^{k}_{ij}}$ は、PyTorchの自動微分で求めます。
次に、二次微分です。
$\displaystyle \frac{\partial^{2} Y^{c}}{(\partial A^{k}_{ij})^{2}} = exp(S^{c}) \left[ \left( \frac{\partial S^{c}}{\partial A^{k}_{ij}} \right)^{2} + \frac{\partial^{2} S^{c}}{(\partial A^{k}_{ij})^{2}} \right] \cdots (14)$
二次微分 $\frac{\partial^{2} S^{c}}{(\partial A^{k}_{ij})^{2}}$ が登場しますが、もし最後の特徴マップから出力層の計算に至るまでにReLUなどの線形変換しか存在していない場合、一次微分は定数になり、二次微分以上は0になります。したがって、(15)式のように $S^{c}$ の一次微分のみで求められる簡単な式にできます。
$\displaystyle \frac{\partial^{2} Y^{c}}{(\partial A^{k}_{ij})^{2}} = exp(S^{c}) \left( \frac{\partial S^{c}}{\partial A^{k}_{ij}} \right)^{2} \cdots (15)$
でも、ここで「この $Y^{c}$ の二次微分が求められるなら変換前の $Y^{c}$ の二次微分も0として求められるんじゃないか？」っと思ったんですがどうなんでしょうか…。まあ0だと(12)式の分母が0になってよくないのはわかりますが…。
三次微分も同様に求められます。
$\displaystyle \frac{\partial^{3} Y^{c}}{(\partial A^{k}_{ij})^{3}} = exp(S^{c}) \left( \frac{\partial S^{c}}{\partial A^{k}_{ij}} \right)^{3} \cdots (16)$
(15)式、(16)式により、(12)式の $\alpha^{kc}_{ij}$ が求まります。その $\alpha^{kc}_{ij}$ を用いて、(11)式の $w^{c}_{k}$ を求めます。
最後に、Grad-CAM++のクラス $c$ に対するheatmapの座標 $(i,j)$ の値 $GradCAMpp^{c}_{ij}$ は次の式で表せます。
$\displaystyle GradCAMpp^{c}_{ij} = ReLU \left( \sum_{k} w^{c}_{k} A^{k}_{ij} \right) \cdots (17)$

heatmapの例

各クラス $c$ で、heatmapを作ってみました。どのクラスでもほぼ同じ部分で赤くなっていることがわかります。(実装合っているのでしょうか…)
CAMのheatmapと比べて赤い部分が多くなっています。局所的な特徴をとらえる特徴マップも重視されるようになった結果でしょうか。

f:id:verliezer93764:20210220120523p:plain — 図5: Grad-CAM++による各クラスでのheatmap出力結果

Score-CAM [Wang+, CVPRW2020]

アイデア

Score-CAMでは勾配情報を用いず、特徴マップでマスクした入力画像をモデルに入力した際の対象クラスのスコアを用いることでどの特徴マップが重要であるか判断します。使用する情報は入力画像と各層の特徴マップのみであるため、任意の層の特徴マップでの可視化も可能です。
まずは、特徴マップ $A^{k}$ を入力画像と同サイズまで拡大し、入力画像のマスクのために特徴マップ $A^{k}$ を $[0, 1]$ に正規化します。max, minはそれぞれ最大輝度値、最小輝度値をとる処理です。
$\displaystyle s(A^{k}) = \frac{A^{k} - min(A^{k})}{max(A^{k}) - min(A^{k})} \cdots (18)$
次に、入力画像 $X_{0}$ と各特徴マップ $A^{k}$ とのアダマール積を求めてマスク済み画像 $M^{k}$ を求めます。
$\displaystyle M^{k} = X_{0} \circ s(A^{k}) \cdots (19)$
次に、マスク済み画像 $M^{k}$ をCNNモデル $F$ に入力して、その結果をさらにsoftmax関数により確率分布へ変換した $S^{k}$ を得ます。
$\displaystyle S^{k} = softmax \left( F(M^{k}) \right) \cdots (20)$
この $S^{k}$ から対象クラス $c$ にあたる出力を抽出し、特徴マップの重要度 $\alpha^{c}_{k}$ を獲得します。
これを利用し、クラス $c$ に関するScore-CAMのheatmapである $ScoreCAM^{c}$ を求め、正規化を行います。
$\displaystyle ScoreCAM^{c} = ReLU \left( \sum_{k} \alpha^{c}_{k} A^{k} \right) \cdots (21)$
$\displaystyle ScoreCAM^{c} \leftarrow \frac{ ScoreCAM^{c} - min \left( ScoreCAM^{c} \right) }{ max \left( ScoreCAM^{c} \right) - min \left( ScoreCAM^{c} \right) } \cdots (22)$

heatmapの例

各クラス $c$ で、heatmapを作ってみました。Grad-CAM++のheatmapと似ています。(これまた実装合っているのでしょうか…)

f:id:verliezer93764:20210220120954p:plain — 図6: Score-CAMによる各クラスでのheatmap出力結果

参考サイトなど

Understanding Pytorch hooks | Kaggle

PyTorchのhook機能のわかりやすい解説。

CVPR 2016 Open Access Repository

CAMの論文。

ICCV 2017 Open Access Repository

Grad-CAMの論文。

Grad-CAM++: Generalized Gradient-Based Visual Explanations for Deep Convolutional Networks | IEEE Conference Publication | IEEE Xplore

Grad-CAM++の論文。

CVPR 2020 Open Access Repository

Score-CAMの論文。

Grad-CAM、Grad-CAM++、Score-CAMを実装・比較してみた – 戦略コンサルで働くデータサイエンティストのブログ

Grad-CAM, Grad-CAM++, Score-CAMの紹介とKerasを使った実装。

GitHub - ramprs/grad-cam: [ICCV 2017] Torch code for Grad-CAM

Grad-CAMの著者らによる実装。

GitHub - adityac94/Grad_CAM_plus_plus: A generalized gradient-based CNN visualization technique

Grad-CAM++の著者らによる実装。

最新手法！「Grad-CAM++」のレビューと実装 - Qiita

Grad-CAM++の解説記事。

GitHub - haofanwang/Score-CAM: [CVPRW 2020] Official implementation of Score-CAM in Pytorch

Score-CAMの著者らによる実装。

[Zhou+, CVPR2016] B. Zhou, A. Khosla, A. Lapedriza, A. Oliva and A. Torralba, "Learning Deep Features for Discriminative Localization," in Proc. IEEE Conference on Computer Vision and Pattern Recognition, 2016, pp. 2921-2929.
[Selvaraju+, ICCV2017] R. R. Selvaraju, M. Cogswell, A. Das, R. Vedantam, D. Parikh, and D. Batra, "Grad-cam: Visual explanations from deep networks via gradient-based localization," in Proc. IEEE International Conference on Computer Vision, 2017, pp. 618-626.
[Chattopadhay+, WACV2018] A. Chattopadhay, A. Sarkar, P. Howlader, and V. N. Balasubramanian, "Grad-cam++: Generalized gradient-based visual explanations for deep convolutional networks," in Proc. IEEE Winter Conference on Applications of Computer Vision, 2018, pp. 839-847.
[Wang+, CVPRW2020] H. Wang, Z. Wang, M. Du, F. Yang, Z. Zhang, S. Ding, P. Mardziel, and X. Hu, "Score-cam: Score-weighted visual explanations for convolutional neural networks," in Proc. IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, 2020, pp. 24-25.