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

时间:2019-06-18 22:48:23

标签: python tensorflow machine-learning keras

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

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

我已经看到了(相当有限的)guide,它在keras中创建自定义图层,但是它没有描述如果我们希望我们的操作具有自定义渐变应该怎么做。

1 个答案:

答案 0 :(得分:4)

首先,在keras下对API(如您所称)的“统一”并不能阻止您像在TensorFlow 1.x中那样做。会话可能不见了,但您仍然可以像定义任何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

请注意,在函数中,您应该将xdy视为张量,而 numpy数组(即执行张量运算)

b)创建一个执行您的custom_op的自定义keras图层。对于此示例,我假设您的图层没有任何可训练的参数或更改其输入的形状,但是如果有的话,它并没有太大的区别。为此,您可以参考发布的支票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