отдельные градиенты по тензорной размерности в тензорном потоке / керас

Question

Я пытаюсь создать фильтр визуализаций, как из этого проекта https://github.com/tensorflow/lucid/tree/master/lucid/optvis

Мне удалось заставить работать какой-то наивный код, который генерирует одно изображение за раз

import time
import os
import multiprocessing as mp
import numpy as np
import tensorflow as tf
from tensorflow.keras.applications.vgg16 import preprocess_input
import matplotlib.pyplot as plt
from PIL import Image
os.environ["CUDA_VISIBLE_DEVICES"]="1"
# Layer name to inspect
layer_name = 'block5_conv2'
temp = tf.zeros([4, 32, 32, 3])  # Or tf.zeros
preprocess_input(temp)
epochs = 1500
step_size = .05
# filter_index = 1

# Create a connection between the input and the target layer
model = tf.keras.models.load_model('base_model_cifar10_vgg16.h5')


# Initiate random noise
input_img_data = np.random.random((1, 224, 224, 3))
input_img_data = (input_img_data - 0.5) * 20 + 128.

# Cast random noise from np.float64 to tf.float32 Variable

targets = []
for i, layer in enumerate(model.layers):
    if 'conv' in layer.name:
        targets.append(layer.name)
targets

%%time
for layer_name in targets:
    print(f'generating dataset {layer_name}')
    submodel = tf.keras.models.Model([model.inputs[0]], 
                                     [model.get_layer(layer_name).output])
    for filter_index in range(submodel.output.get_shape()[-1]):
        input_img_data = np.random.random((1, 224, 224, 3))
        input_img_data = (input_img_data - 0.5) * 20 + 128.
        input_img_data = tf.Variable(tf.cast(input_img_data, tf.float32))
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                outputs = submodel(preprocess_input(input_img_data))
                loss_value = tf.reduce_mean(outputs[:, :, :, filter_index])
            grads = tape.gradient(loss_value, input_img_data)
            normalized_grads = grads / (tf.sqrt(tf.reduce_mean(tf.square(grads))) + 1e-5)
            input_img_data.assign_add(normalized_grads * step_size)
            image = input_img_data.numpy()[0,:,:,:].astype(np.uint8)
        p = mp.Process(target=Image.fromarray(image).save(fp=f'./synth_sets/vgg16_cifar10/{layer_name}_{filter_index}.jpg'))
        p.start()

, но ДЕЙСТВИТЕЛЬНО медленно, поскольку генерирует одно изображение за раз. Это похоже на большую трату циклов графического процессора.

Тогда у меня возникла идея попробовать одновременно создать серию изображений с тем же размером, что и количество фильтров в слое.

Только я не могу понять, как сгенерировать градиент по отношению к каждому срезу тензора на оси изображения

Это уже довольно далеко, и я получил только список типов None.

for layer_name in targets[0:1]:
    print(f'generating dataset {layer_name}')
    submodel = tf.keras.models.Model([model.inputs[0]], 
                                     [model.get_layer(layer_name).output])

    batch = submodel.output.get_shape()[-1]
    input_img_data = np.random.random((batch, 224, 224, 3))
    input_img_data = (input_img_data - 0.5) * 20 + 128.
    input_img_data = tf.Variable(tf.cast(input_img_data, tf.float32))
    input_img_list = tf.split(input_img_data, batch, axis=0)
    zero_init = tf.zeros_initializer()
    loss_value = [tf.Variable(0, dtype=tf.float32) for i in range(batch)]
    #print(loss_value)
    for _ in range(epochs):
        #print(tf.split(input_img_data, batch, axis=0))
        with tf.GradientTape() as tape:
            outputs = tf.dtypes.cast(submodel(preprocess_input(input_img_data)), tf.float32)
            for i in range(batch):

                ph = tf.Variable(initial_value=zero_init(shape=(224, 224, 3)),dtype=tf.float32)
                ph = tf.reduce_mean(outputs[:,:,:,i])
                #print(loss_value)
                loss_value[i].assign(ph)

            grads = tape.gradient(loss_value, tf.split(input_img_data, batch, axis=0))
        print(grads)
#         normalized_grads = grads / (tf.sqrt(tf.reduce_mean(tf.square(grads))) + 1e-5)
#         input_img_data.assign_add(normalized_grads * step_size)

    for i in range(batch):
        image = input_img_data.numpy()[i,:,:,:].astype(np.uint8)
        p = mp.Process(target=Image.fromarray(image).save(fp=f'./synth_sets/vgg16_cifar10_test/{layer_name}_{i}.jpg'))
        p.start()

Кто-нибудь знает, как я могу рассчитать индивидуальный градиент для среза изображения?