kerasでのGrad-CAMの実装でエラーが発生する。

Question

前提・実現したい事

深層距離学習用に構築した学習モデルでgrad-camの実装を行いたいのですが、
エラーが発生し、調べても自己解決できていない状況です。
ご存じの方おりましたら、エラーの原因と対応の方法を教えて頂けますでしょうか。
参考にした記事は下記のとおりです。

学習モデルの構築
https://qiita.com/noritsugu_yamada/items/2e049cd7a8fd77eee0f5#_reference-848bd3284d2b1b43f22d

https://qiita.com/daisukelab/items/e0ff429bd58b2befbb1b

Grad-CAMの実装
https://qiita.com/Kensuke-N/items/cea3920f19bebc9e0d19

発生している問題・エラーメッセージ

下記コマンド実行時にエラーが発生しました。

cam = grad_cam(model, x, target_layer)

発生したエラーは以下の通りです。

AssertionError: Could not compute output Tensor("activation/Softmax_2:0", shape=(None, 2), dtype=float32)

コード

import tensorflow as tf
import keras
import numpy as np
import cv2
from keras import backend as K
# 画像用
from keras.preprocessing.image import array_to_img, img_to_array, load_img
# モデル読み込み用
from keras.models import load_model
# Grad−CAM計算用
from tensorflow.keras import models
import tensorflow as tf
from keras.engine.topology import Layer

class Arcfacelayer(Layer):
    # s:softmaxの温度パラメータ, m:margin
    def __init__(self, output_dim, s=30, m=0.50, easy_magin=False, **kwargs):
        self.output_dim = output_dim
        self.s = s
        self.m = m
        self.easy_magin = easy_magin
        super(Arcfacelayer, self).__init__()
        
########### モデルの保存用に追加  ##########################   
    def get_config(self):
        config = {
            "output_dim" : self.output_dim,
            "s" : self.s,
            "m" : self.m,
            "easy_magin" : self.easy_magin,
        }
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))
##########################################################

# 重みの作成
    def build(self, input_shape):
        # Create a trainable weight variable for this layer.
        self.kernel = self.add_weight(name='kernel',
                                      shape=(input_shape[0][1], self.output_dim),
                                      initializer='uniform',
                                      trainable=True)
        super(Arcfacelayer, self).build(input_shape)

# mainの処理
    def call(self, x):
        y = x[1]
        x_normalize = tf.math.l2_normalize(x[0]) #|x = x'/ ||x'||2
        k_normalize = tf.math.l2_normalize(self.kernel) # Wj = Wj' / ||Wj'||2
        
        cos_m = K.cos(self.m)
        sin_m = K.sin(self.m)
        th = K.cos(np.pi - self.m)
        mm = K.sin(np.pi - self.m) * self.m
        
        cosine = K.dot(x_normalize, k_normalize) # W.Txの内積
        
        sine = K.sqrt(1.0 - K.square(cosine))
        phi = cosine * cos_m - sine * sin_m
        
        if self.easy_magin:
            phi = tf.where(cosine > 0, phi, cosine) 
            
        else:
            phi = tf.where(cosine > th, phi, cosine - mm) 
        
        # 正解クラス:cos(θ+m) 他のクラス:cosθ 
        output = (y * phi) + ((1.0 - y) * cosine) # true cos(θ+m), False cos(θ)
        output *= self.s
        
        return output
    
    def compute_output_shape(self, input_shape):
        
        return (input_shape[0][0], self.output_dim)
        #return self.output_dim

def grad_cam(input_model, x, layer_name):
    """
    Args: 
        input_model(object): モデルオブジェクト
        x(ndarray): 画像
        layer_name(string): 畳み込み層の名前
    Returns:
        output_image(ndarray): 元の画像に色付けした画像
    """

    # 画像の前処理
    # 読み込む画像が1枚なため、次元を増やしておかないとmode.predictが出来ない
    X = np.expand_dims(x, axis=0)
    preprocessed_input = X.astype('float32') / 255.0
    
    grad_model = models.Model([input_model.inputs], [input_model.get_layer(layer_name).output, input_model.output])

    with tf.GradientTape() as tape:
        conv_outputs, predictions = grad_model(preprocessed_input)
        class_idx = np.argmax(predictions[0])
        loss = predictions[:, class_idx]

    # 勾配を計算
    output = conv_outputs[0]
    grads = tape.gradient(loss, conv_outputs)[0]

    gate_f = tf.cast(output > 0, 'float32')
    gate_r = tf.cast(grads > 0, 'float32')

    guided_grads = gate_f * gate_r * grads

    # 重みを平均化して、レイヤーの出力に乗じる
    weights = np.mean(guided_grads, axis=(0, 1))
    cam = np.dot(output, weights)

    # 画像を元画像と同じ大きさにスケーリング
    cam = cv2.resize(cam, IMAGE_SIZE, cv2.INTER_LINEAR)
    # ReLUの代わり
    cam  = np.maximum(cam, 0)
    # ヒートマップを計算
    heatmap = cam / cam.max()

    # モノクロ画像に疑似的に色をつける
    jet_cam = cv2.applyColorMap(np.uint8(255.0*heatmap), cv2.COLORMAP_JET)
    # RGBに変換
    rgb_cam = cv2.cvtColor(jet_cam, cv2.COLOR_BGR2RGB)
    # もとの画像に合成
    output_image = (np.float32(rgb_cam) + x / 2)  

    return output_image

model_path = 'model_2/m02_fine157kara.hdf5'
image_path = 'dataset/03_test/NG/10.tif'
IMAGE_SIZE  = (224,224)

model = load_model(model_path, custom_objects = {"Arcfacelayer": Arcfacelayer})
model.summary()
x = img_to_array(load_img(image_path, target_size=IMAGE_SIZE))

target_layer = 'out_relu'
cam = grad_cam(model, x, target_layer)

補足情報

python version: 3.7.0
tensorflow version: 2.3.1
keras version: 2.4.3

Comments

[pipi_m3tro]

学習モデルの入力部と出力部は以下の通りです。

Model: "functional_1"
__________________________________________________________________________________________________
Layer (type) Output Shape Param # Connected to
==================================================================================================
input_1 (InputLayer) [(None, 224, 224, 3) 0
__________________________________________________________________________________________________
Conv1_pad (ZeroPadding2D) (None, 225, 225, 3) 0 input_1[0][0]
__________________________________________________________________________________________________
Conv1 (Conv2D) (None, 112, 112, 32) 864 Conv1_pad[0][0]
__________________________________________________________________________________________________
bn_Conv1 (BatchNormalization) (None, 112, 112, 32) 128 Conv1[0][0]
__________________________________________________________________________________________________
Conv1_relu (ReLU) (None, 112, 112, 32) 0 bn_Conv1[0][0]
___________________________________________________________________________________
・
・
・
・
・
__________________________________________________________________________________________________
block_16_project (Conv2D) (None, 7, 7, 320) 307200 block_16_depthwise_relu[0][0]
__________________________________________________________________________________________________
block_16_project_BN (BatchNorma (None, 7, 7, 320) 1280 block_16_project[0][0]
__________________________________________________________________________________________________
Conv_1 (Conv2D) (None, 7, 7, 1280) 409600 block_16_project_BN[0][0]
__________________________________________________________________________________________________
Conv_1_bn (BatchNormalization) (None, 7, 7, 1280) 5120 Conv_1[0][0]
__________________________________________________________________________________________________
out_relu (ReLU) (None, 7, 7, 1280) 0 Conv_1_bn[0][0]
__________________________________________________________________________________________________
global_average_pooling2d (Globa (None, 1280) 0 out_relu[0][0]
__________________________________________________________________________________________________
input_2 (InputLayer) [(None, 2)] 0
__________________________________________________________________________________________________
arcfacelayer (Arcfacelayer) (None, 2) 2560 global_average_pooling2d[0][0]
input_2[0][0]
__________________________________________________________________________________________________
activation (Activation) (None, 2) 0 arcfacelayer[0][0]
==================================================================================================
Total params: 2,260,544
Trainable params: 2,226,432
Non-trainable params: 34,112
________________________________________

[退会済みユーザー]

たぶん以下の記事のコードですよね？これからどこか変更しましたか？
https://qiita.com/noritsugu_yamada/items/2e049cd7a8fd77eee0f5

後は似たような現象の報告があったので参考にURLを貼り付けます。https://stackoverflow.com/questions/61997378/assertionerror-could-not-compute-output-tensor

[pipi_m3tro]

ご回答ありがとうございます。
記事のコードから変更した点は、学習時にモデル全体を保存する必要があった為、Arcfacelayerクラスにget_config関数を追加したことと
Grad-CAM実装時のmodel = load_model()内に custom_objects = {"Arcfacelayer": Arcfacelayer}を追加したことになります。

参考URL確認しましたが、今ケースの修正箇所が分かりませんでした。。

[退会済みユーザー]

Arcfaceでgradcamしたい、が目的であれば以下のURLのコードであれば動かせるかもしれません。
https://ni4muraano.hatenablog.com/entry/2019/08/14/111620
# CNN関連バグ取りはめんどくさいですよね💦