python - 如何在TF2.0中创建自定义渐变的keras图层？

Question

因为在 TensorFlow 2.0 中他们计划统一 keras(我不太熟悉)下的所有高级 API 并完全删除 session ，我想知道:

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

我看过(相当有限的)guide关于在 keras 中创建自定义层，但它没有描述如果我们希望我们的操作具有自定义梯度，我们应该做什么。

Answer 1

首先，keras 下 API(如您所说)的“统一”不会阻止您像在 TensorFlow 1.x 中那样做事。 session 可能已经消失，但您仍然可以像任何 python 函数一样定义您的模型，并在没有 keras 的情况下热切地训练它(即通过 tf.GradientTape )

现在，如果您想构建一个带有执行自定义操作并具有自定义梯度的自定义层的keras模型，您应执行以下操作:

a) 编写一个函数来执行您的自定义操作并定义您的自定义梯度。有关如何执行此操作的更多信息 here .

@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 视为 Tensors 而不是 numpy 数组(即执行张量运算)

b) 创建一个自定义 keras 层来执行您的custom_op。对于这个例子，我假设你的层没有任何可训练的参数或改变它的输入的形状，但如果有的话也没有太大区别。为此，您可以引用发布支票的指南 this one .

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