Keras GradCam 实现，可以处理成批图像而不是一次处理单个图像

Question

我正在关注 Keras 文档https://keras.io/examples/vision/grad_cam/中的 GradCam 示例，并希望对其进行修改，以便它可以处理一批图像，而不是一次只能处理一个图像。

我已经能够做到这一点，但不得不使用对tf.map_fn的调用，我想摆脱它以期提高性能。

到目前为止我的进展（ Google Coolab的完整代码）：

#https://keras.io/examples/vision/grad_cam/
def make_gradcam_heatmap(grad_model, images, pred_index=None):
    images = tf.cast(images, tf.float32)

    # Then, we compute the gradient of the top predicted class for our input image
    # with respect to the activations of the last conv layer
    with tf.GradientTape() as tape:
        tape.watch(images)
        last_conv_layer_output, preds = grad_model(images)
        if pred_index is None:
            pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]

    # This is the gradient of the output neuron (top predicted or chosen)
    # with regard to the output feature map of the last conv layer
    grads = tape.gradient(class_channel, last_conv_layer_output)
    assert grads is not None, "GradientTape returned gradients=None"

    # This is a vector where each entry is the mean intensity of the gradient
    # over a specific feature map channel
    pooled_grads = tf.reduce_mean(grads, axis=(1, 2))

    # We multiply each channel in the feature map array
    # by "how important this channel is" with regard to the top predicted class
    # then sum all the channels to obtain the heatmap class activation
    def single_image(index):
        return last_conv_layer_output[index] @ pooled_grads[index][tf.newaxis, ..., tf.newaxis]

    heatmaps = tf.map_fn(single_image, tf.range(tf.shape(grads)[0]), dtype=tf.float32)

    # normalize the whole batch to [0, 1]
    heatmaps -= tf.math.reduce_min(heatmaps, axis=(0,1,2,3))
    heatmaps /= tf.math.reduce_max(heatmaps, axis=(0,1,2,3))

    return heatmaps

---

有没有什么办法可以不使用tf.map_fn来重写这段代码？

    def single_image(index):
        return last_conv_layer_output[index] @ pooled_grads[index][tf.newaxis, ..., tf.newaxis]

heatmaps = tf.map_fn(single_image, tf.range(tf.shape(grads)[0]), dtype=tf.float32)

Answer 1

你可以使用matmul ，

heatmaps = tf.matmul(last_conv_layer_output, pooled_grads[:,tf.newaxis, :, tf.newaxis])