如何使用 tensorflow 1.14 为 keras 层定义自定义渐变

Question

我正在尝试在 tf.keras（tf 版本 1.14）中实现渐变反转层。我使用 tf.custom_gradient 定义了一个自定义渐变，但从未调用过 grad 函数，结果是错误的。

我正在使用以下来测试图层

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras import backend as k
import tensorflow as tf
tf.enable_eager_execution()
import numpy as np

@tf.custom_gradient
def custom_op(x):
    result = x
    def custom_grad(dy):
        grad = tf.negative(dy)
        return grad
    return result, custom_grad

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

  def call(self, inputs):
    return custom_op(inputs)


a = Input(shape=(2,))
b = GradientReversal()(a)
c = Dense(1)(b)
model = Model(inputs=a, outputs=c)
model.compile(loss='binary_crossentropy', optimizer='sgd')

@tf.function
def test():  
  with tf.GradientTape() as tape:
    out = model(np.ones((1, 2)))
  gradients = tape.gradient(out, model.trainable_variables)
  return gradients

print(test())

Answer 1

没关系，我应该在 Dense 之后使用 GradientReversal 层来获得我想要的结果。