Tensorflow: создайте собственный слой Conv2d с настраиваемым градиентом

Question

Я хотел бы иметь слой Conv2d, который может возвращать градиент со значениями по определенным индексам, умноженными на 0,2. Первоначальная функциональность должна остаться.

Моя идея состояла в том, чтобы создать подкласс класса Conv2d и иметь метод установки для сохранения индексов, где Groundtruth равен 0:

import tensorflow as tf

class BackpassFilter(tf.keras.layers.Layer):
    def __init__(self, kernel_initializer, num_filter=48, kernel_size=3, stride=1, padding='same', name='keypoints_3_4'):
    super(BackpassFilter, self).__init__()
    self.fc = tf.keras.layers.Conv2D(kernel_initializer, num_filter, kernel_size, stride, padding, name, kernel_initializer=kernel_initializer)

    def call(self, input):
        return self.fc(input)

    def add_filter_indices(self, ground_truth):
        self.tmp_indices = tf.where(tf.equal(ground_truth, 0))

    @tf.custom_gradient
    def custom_gradient(self, x):
        ...

Я обнаружил, что @tf .custom_gradient - это то, что я мог бы использовать. Как я могу реализовать метод доступа к исходным градиентам, управления ими и их возврата? Должен ли я как-то регистрировать настраиваемый градиент, чтобы использовать его в классе?