TF 2.0如何实现梯度反转层？

Question

这个层是静态的，是一个伪函数。在前向传播中它什么都不做（恒等函数）。然而，在反向传播中，它将梯度乘以 -1。 github 上有很多实现，但它们不适用于 TF 2.0。

这里有一个供参考。

import tensorflow as tf
from tensorflow.python.framework import ops

class FlipGradientBuilder(object):
    def __init__(self):
        self.num_calls = 0

    def __call__(self, x, l=1.0):
        grad_name = "FlipGradient%d" % self.num_calls
        @ops.RegisterGradient(grad_name)
        def _flip_gradients(op, grad):
            return [tf.negative(grad) * l]

        g = tf.get_default_graph()
        with g.gradient_override_map({"Identity": grad_name}):
            y = tf.identity(x)

        self.num_calls += 1
        return y

flip_gradient = FlipGradientBuilder()

Answer 1

反转梯度的虚拟操作

这可以使用装饰器 tf.custom_gradient 完成，如 in this example:

所述

@tf.custom_gradient
def grad_reverse(x):
    y = tf.identity(x)
    def custom_grad(dy):
        return -dy
    return y, custom_grad

然后，您就可以像使用普通的 TensorFlow op 一样使用它，例如：

z = encoder(x)
r = grad_reverse(z)
y = decoder(r)

凯拉斯 API?

TF 2.0 的一大便利是它原生支持 Keras API。您可以定义自定义 GradReverse 操作并享受 Keras 的便利：

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

    def call(self, x):
        return grad_reverse(x)

然后，您可以将此层用作 Keras 的任何其他层，例如：

model = Sequential()
conv = tf.keras.layers.Conv2D(...)(inp)
cust = CustomLayer()(conv)
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=...)
model.fit(...)

TF 2.0如何实现梯度反转层？

How to implement gradient reversal layer in TF 2.0?

keras

tensorflow

tensorflow2.0

反转梯度的虚拟操作

凯拉斯 API?