如何在TF2.0中使用自定义渐变创建keras图层？

Question

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

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

我看过（非常有限的）关于在keras中创建自定义图层的指南，但是它没有描述如果我们希望我们的操作具有自定义渐变，我们应该怎么做。

Answer 1

首先，在keras下对API（如您所称的）的“统一”并不会阻止您像在TensorFlow 1.x中那样做。会话可能消失了，但是您仍然可以像定义任何python函数一样定义模型，并在没有keras的情况下积极地训练它（即通过tf.GradientTape）

现在，如果要使用执行自定义操作并具有自定义渐变的自定义层构建keras模型，则应执行以下操作：

a）编写执行自定义操作并定义自定义渐变的函数。有关如何执行此操作的更多信息，请参见此处。

@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和dy视为张量，而不是 numpy数组（即执行张量运算）

b）创建一个执行您的的定制keras图层custom_op。对于此示例，我假设您的图层没有任何可训练的参数或更改其输入的形状，但是如果有的话，它并没有多大区别。为此，您可以参考发布的指南，检查该指南。

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