如何在TF2.0中创建具有自定义渐变的keras层？

由于在 TensorFlow 2.0 中，他们计划统一 keras 下的所有高级 API（我不太熟悉）并完全删除会话，我想知道：

如何创建具有自定义渐变的自定义 keras 层？

我见过（相当有限）guide https://keras.io/layers/writing-your-own-keras-layers/关于在 keras 中创建自定义图层，但它没有描述如果我们希望我们的操作具有自定义渐变我们应该做什么。

首先，keras 下 API（如您所称）的“统一”并不会阻止您像在 TensorFlow 1.x 中那样进行操作。会话可能会消失，但您仍然可以像任何 python 函数一样定义模型，并在没有 keras 的情况下热切地训练它（即通过tf.渐变带 https://www.tensorflow.org/tutorials/eager/custom_training_walkthrough)

现在，如果你想构建一个 keras 模型自定义层执行一个自定义操作并有一个自定义渐变，您应该执行以下操作：

a) 编写一个执行自定义操作并定义自定义渐变的函数。有关如何执行此操作的更多信息here https://www.tensorflow.org/api_docs/python/tf/custom_gradient.

@tf.custom_gradient
def custom_op(x):
    result = ... # do forward computation
    def custom_grad(dy):
        grad = ... # compute gradient
        return grad
    return result, custom_grad

请注意，在函数中您应该对待x and dy作为张量和notnumpy 数组（即执行张量运算）

b) 创建一个自定义 keras 层来执行您的custom_op。对于这个例子，我假设你的层没有任何可训练的参数或改变其输入的形状，但如果有的话也没有太大区别。为此，您可以参考您发布的检查指南this one https://www.tensorflow.org/beta/tutorials/eager/custom_layers.

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(CustomLayer, self).__init__()

    def call(self, x):
        return custom_op(x)  # you don't need to explicitly define the custom gradient
                             # as long as you registered it with the previous method

现在您可以在 keras 模型中使用该层并且它会起作用。例如：

inp = tf.keras.layers.Input(input_shape)
conv = tf.keras.layers.Conv2D(...)(inp)  # add params like the number of filters
cust = CustomLayer()(conv)  # no parameters in custom layer
flat = tf.keras.layers.Flatten()(cust)
fc = tf.keras.layers.Dense(num_classes)(flat)

model = tf.keras.models.Model(inputs=[inp], outputs=[fc])
model.compile(loss=..., optimizer=...)  # add loss function and optimizer
model.fit(...)  # fit the model