如何使用预先计算的梯度来更新神经网络权重

Question

我正在研究一个基于物理的神经网络问题，我需要用 NN 输出 y 的物理评估来替换损失函数。这个评估给了我一个梯度来更新y，我想用这个梯度来更新NN的权重。我不喜欢特定的框架，只要它有效。

我尝试使用以下代码使用张量流：

@tf.custom_gradient
def fake_loss(x):
    def grad(dy):
        # grad=gradtrue
        return gradtrue*dy
    return x,grad

for e in range(10):
  with tf.GradientTape() as tape:
    pred = model(x)
    gradtrue=tf.convert_to_tensor(physics_gradient(x))
    loss = fake_loss(pred)
  grads = tape.gradient(loss, model.trainable_weights)
  optimizer.apply_gradients(zip(grads, model.trainable_weights))

其中 physics_gradient 函数是提到的 NN 输出的物理评估。 这会编译，但不会减少 physics_gradient

给出的损失误差