且构网

分享程序员开发的那些事...
且构网 - 分享程序员编程开发的那些事

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

更新时间:2023-11-17 23:31:10

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

First of all, the "unification" of the APIs (as you call it) under keras doesn't prevent you from doing things like you did in TensorFlow 1.x. Sessions might be gone but you can still define your model like any python function and train it eagerly without keras (i.e. through tf.GradientTape)

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

Now, if you want to build a keras model with a custom layer that performs a custom operation and has a custom gradient, you should do the following:

a)编写执行自定义操作并定义自定义渐变的函数. 此处.

a) Write a function that performs your custom operation and define your custom gradient. More info on how to do this 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数组(即执行张量运算)

Note that in the function you should treat x and dy as Tensors and not numpy arrays (i.e. perform tensor operations)

b)创建一个执行您的custom_op的自定义keras图层.对于此示例,我将假设您的图层没有任何可训练的参数或更改其输入的形状,但是如果有的话,它并没有多大区别.为此,您可以参考发布的指南,请选中.. >

b) Create a custom keras layer that performs your custom_op. For this example I'll assume that your layer doesn't have any trainable parameters or change the shape of its input, but it doesn't make much difference if it does. For that you can refer to the guide that you posted check 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模型中使用此层,它将起作用.例如:

Now you can use this layer in a keras model and it will work. For example:

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