在Tensorflow GradientTape中声明变量

Question

我有一个模型，有一个复杂的损失，计算每类模型的输出。

如下所示，我使用一些自定义的损失函数计算损失，将这个值赋值给变量，因为张量在tensorflow中是不变的。

def calc_loss(y_true, y_pred):
    num_classes=10
    pos_loss_class = tf.Variable(tf.zeros((1, num_classes), dtype=tf.dtypes.float32))
    for idx in range(num_classes):
        pos_loss = SOME_LOSS_FUNC(y_true[:, idx], y_pred[:, idx] 
        pos_loss_class[:, idx].assign(pos_loss)

return tf.reduce_mean(pos_loss_class)

我的代码很简单：

with tf.GradientTape() as tape:
    output = model(input, training=True)
    loss = calc_loss(targets, output)
grads = tape.gradient(loss, model.trainable_weights)

但是，对于所有模型的变量，我都没有收到任何信息。根据我的理解，这是由如下所写的变量状态的阻塞方式造成的：https://www.tensorflow.org/guide/autodiff#4_took_gradients_through_a_stateful_object

有什么建议吗？

下面是可重复的代码，这是一个玩具示例，但演示了这个问题。

y_true = tf.Variable(tf.random.normal((1, 2)), name='targets')
layer = tf.keras.layers.Dense(2, activation='relu')
x = tf.constant([[1., 2., 3.]])
with tf.GradientTape() as tape:
    y_pred = layer(x)
    loss_class = tf.Variable(tf.zeros((1,2)), dtype=tf.float32)
    for idx in range(2):
        loss = tf.abs(y_true[:, idx] - y_pred[:, idx])
        loss_class[:, idx].assign(loss)
    final_loss = tf.reduce_mean(loss_class)
grads = tape.gradient(final_loss, layer.trainable_weights)

Answer 1

我目前的第二个猜测是，赋值方法阻止梯度，正如您喜欢的tensorflow页面中所解释的那样.相反，只使用一个简单的列表：

def calc_loss(y_true, y_pred):
    num_classes=10
    pos_loss_class = []
    for idx in range(num_classes):
        pos_loss = SOME_LOSS_FUNC(y_true[:, idx], y_pred[:, idx] 
        pos_loss_class.append(pos_loss)

    return tf.reduce_mean(pos_loss_class)