在tf 2.0中计算GradCam的梯度

import tensorflow as tf
import numpy as np
import cv2

model = tf.keras.applications.VGG16(weights='imagenet', include_top=True)

def grad_cam(input_model, image, layer_name):
    # 获取指定层的输出和梯度
    grad_model = tf.keras.models.Model([input_model.inputs], [input_model.get_layer(layer_name).output, input_model.output])
    with tf.GradientTape() as tape:
        conv_output, predictions = grad_model(np.array([image]))
        loss = predictions[:, np.argmax(predictions[0])]
    grads = tape.gradient(loss, conv_output)[0]

    # 计算通道权重
    weights = np.mean(grads, axis=(0, 1))
    cam = np.dot(conv_output[0], weights)

    # 反向传播
    cam = cv2.resize(cam, (image.shape[1], image.shape[0]))
    cam = np.maximum(cam, 0)
    heatmap = cam / np.max(cam)

    return heatmap

image_path = 'path_to_image.jpg'
image = cv2.imread(image_path)
image = cv2.resize(image, (224, 224))
image = image / 255.0

heatmap = grad_cam(model, image, 'block5_conv3')
heatmap = cv2.applyColorMap(np.uint8(255 * heatmap), cv2.COLORMAP_JET)
superimposed_img = cv2.addWeighted(image, 0.6, heatmap, 0.4, 0)

cv2.imshow('GradCam', superimposed_img)
cv2.waitKey(0)
cv2.destroyAllWindows()