在 TensorFlow/Keras 中，如何在自定义 RNN 单元内使用“add_loss”方法？

Question

我的目标：在自定义 RNN 单元内（在图形执行模式下）使用 add_loss 方法添加依赖于输入的损失。

一般设置：

• 使用 Python 3.9
• 使用 TensorFlow 2.8 或 2.10
• 假设import tensorflow as tf，我有一个子类tf.keras.Model，它使用标准tf.keras.layers.RNN层和自定义RNN单元（子类tf.keras.layers.Layer）。在我的自定义 RNN 单元内，我调用 self.add_loss(*) 以添加依赖于输入的损失。

预期结果: 当我调用Model.fit() 时，add_loss 方法会为每个批次和每个时间步调用。梯度计算步骤使用增加的损失而不会引发错误。

实际结果：当我调用 Model.fit() 时，在梯度计算步骤中会引发 InaccessibleTensorError，特别是在 Model.train_step() 中调用 self.losses 时。

Exception has occurred: InaccessibleTensorError
<tf.Tensor 'foo_model/rnn/while/bar_cell/Sum_1:0' shape=() dtype=float32> is out of scope and cannot be used here. Use return values, explicit Python locals or TensorFlow collections to access it.
Please see https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values for more information.

我试过的:

• 错误是不是在使用 unroll=True 初始化 RNN 层时引发（使用 eager- 或 graph-execution）。不幸的是，这对我没有帮助，因为我的序列可能很长。在调试时检查 self.losses 显示正确数量的元素（即 4，每个时间步一个）。
• 错误是不是使用 eager execution 和 unroll=False 时引发。但是检查 self.losses 显示 self.losses 中的元素数量不正确；有一个额外的元素（即 5）。进一步调查显示，有一个额外的呼叫add_loss。不知道为什么会这样。
• 切换到最新稳定版本的 TensorFlow (2.10.0) 无法解决问题。
• 在搜索网络、Stack Overflow 和 TensorFlow 的 GitHub 上的问题/代码后，我完全被难住了。

最小可重现示例

• 使用pytest <name_of_file>.py 从命令行运行。

import pytest
import tensorflow as tf


class FooModel(tf.keras.Model):
    """A basic model for testing.

    Attributes:
        cell: The RNN cell layer.

    """

    def __init__(self, rnn=None, **kwargs):
        """Initialize.

        Args:
            rnn: A Keras RNN layer.
            kwargs:  Additional key-word arguments.

        Raises:
            ValueError: If arguments are invalid.

        """
        super().__init__(**kwargs)

        # Assign layers.
        self.rnn = rnn

    def call(self, inputs, training=None):
        """Call.

        Args:
            inputs: A dictionary of inputs.
            training (optional): Boolean indicating if training mode.

        """
        output = self.rnn(inputs, training=training)
        return output


class BarCell(tf.keras.layers.Layer):
    """RNN cell for testing."""
    def __init__(self, **kwargs):
        """Initialize.

        Args:

        """
        super(BarCell, self).__init__(**kwargs)

        # Satisfy RNNCell contract.
        self.state_size = [tf.TensorShape([1]),]

    def call(self, inputs, states, training=None):
        """Call."""
        output = tf.reduce_sum(inputs, axis=1) + tf.constant(1.0)
        self.add_loss(tf.reduce_sum(inputs))

        states_tplus1 = [states[0] + 1]
        return output, states_tplus1


@pytest.mark.parametrize(
    "is_eager", [True, False]
)
@pytest.mark.parametrize(
    "unroll", [True, False]
)
def test_rnn_fit_with_add_loss(is_eager, unroll):
    """Test fit method (triggering backprop)."""
    tf.config.run_functions_eagerly(is_eager)

    # Some dummy input formatted as a TF Dataset.
    n_example = 5
    x = tf.constant([
        [[1, 2, 3], [2, 0, 0], [3, 0, 0], [4, 3, 4]],
        [[1, 13, 8], [2, 0, 0], [3, 0, 0], [4, 13, 8]],
        [[1, 5, 6], [2, 8, 0], [3, 16, 0], [4, 5, 6]],
        [[1, 5, 12], [2, 14, 15], [3, 17, 18], [4, 5, 6]],
        [[1, 5, 6], [2, 14, 15], [3, 17, 18], [4, 5, 6]],
    ], dtype=tf.float32)
    y = tf.constant(
        [
            [[1], [2], [1], [2]],
            [[10], [2], [1], [7]],
            [[4], [2], [6], [2]],
            [[4], [2], [1], [2]],
            [[4], [2], [1], [2]],
        ], dtype=tf.float32
    )
    ds = tf.data.Dataset.from_tensor_slices((x, y))
    ds = ds.batch(n_example, drop_remainder=False)

    # A minimum model to reproduce the issue.
    cell = BarCell()
    rnn = tf.keras.layers.RNN(cell, return_sequences=True, unroll=unroll)
    model = FooModel(rnn=rnn)
    compile_kwargs = {
        'loss': tf.keras.losses.MeanSquaredError(),
        'optimizer': tf.keras.optimizers.Adam(learning_rate=.001),
    }
    model.compile(**compile_kwargs)

    # Call fit which will trigger gradient computations and raise an error
    # during graph execution.
    model.fit(ds, epochs=1)

Answer 1

我可以确认我遇到了同样的问题。评论以增加知名度。另外，我创建了一个与此相关的 github 问题 (https://github.com/tensorflow/tensorflow/issues/59319)。