通过神经网络在 Tensorflow 中反向传播自定义梯度

Question

这是我要解决的问题：我有一个输出给定大小图像的神经网络；然后将这些图像传递给一个在 TF 中不可微分的函数（假设它完全在程序范围之外），并将此输出与标签进行比较以获得损失。不过，此函数还返回其损失相对于最后一层的梯度，因此我试图对其进行反向传播。 从我收集到的信息来看，

@tf.custom_gradient

应该解决这个问题，我只是不确定如何实施它。下面是我想做的事情以及我尝试过的事情的例子（这可能是错误的）

out = Conv2D(1, (1,1), activation='sigmoid') #this is the last layer of the model
model = Model(input, out)

@tf.custom_gradient
def new_op(x): # x is the output of the model
    def grad(upstream): # here it should return the gradient of the function wrt to x
        result = tf.convert_to_tensor(my_complex_function(x)) 
        return upstream * result
    return x, grad

#training loop

for e in range(epochs):
  print(f'epoch {e}')
  with tf.GradientTape() as tape:
    pred = model(x) #x is the input, defined elsewhere
    loss = new_op(vel_uso)
  grads = tape.gradient(loss, model.trainable_weights)
  optimizer.apply_gradients(zip(grads, model.trainable_weights))

现在这在它运行的意义上是有效的，但是检查结果和自定义梯度，它们不太一致。我试图 trick TF 通过使用它作为损失来传播梯度，虽然这不是我想要最小化的损失，但我确实希望这个梯度变得尽可能小。我知道我可能应该使用

Variables

的

tf.custom_gradient

部分，我只是不确定如何添加它以及训练循环的样子。 感谢任何帮助，如果需要我可以尝试澄清