在子类 tf.keras.Model 的调用方法中使用 GRUCell 进行 For 循环

Posted 2023-02-16

【中文标题】在子类 tf.keras.Model 的调用方法中使用 GRUCell 进行 For 循环

【英文标题】：For loop with GRUCell in call method of subclassed tf.keras.Model

【发布时间】：2021-10-10 21:27:27

【问题描述】：

我对@9​​87654322@ 进行了子类化，并在 for 循环中使用tf.keras.layers.GRUCell 来计算序列“y_t”（n，时间步长，hidden_​​units）和最终隐藏状态“h_t”（n，hidden_​​units）。为了让我的循环输出“y_t”，我在每次循环迭代后更新tf.Variable。使用 model(input) 调用模型不是问题，但是当我在调用方法中使用 for 循环拟合模型时，我会得到 TypeError 或 ValueError。

请注意，我不能简单地使用tf.keras.layers.GRU，因为我正在尝试实现这个paper。本文不只是将 x_t 传递给 RNN 中的下一个单元，而是执行一些计算作为 for 循环中的一个步骤（它们在 PyTorch 中实现）并将计算结果传递给 RNN 单元。他们最终基本上是这样做的：h_t = f(special_x_t, h_t-1)。

请查看以下导致错误的型号：

class CustomGruRNN(tf.keras.Model):
    def __init__(self, batch_size, timesteps, hidden_units, features, **kwargs):

        # Inheritance
        super().__init__(**kwargs)

        # Args
        self.batch_size = batch_size
        self.timesteps = timesteps
        self.hidden_units = hidden_units        

        # Stores y_t
        self.rnn_outputs = tf.Variable(tf.zeros(shape=(batch_size, timesteps, hidden_units)), trainable=False)

        # To be used in for loop in call
        self.gru_cell = tf.keras.layers.GRUCell(units=hidden_units)

        # Reshape to match input dimensions
        self.dense = tf.keras.layers.Dense(units=features)

    def call(self, inputs):
        """Inputs is rank-3 tensor of shape (n, timesteps, features) """

        # Initial state for gru cell
        h_t = tf.zeros(shape=(self.batch_size, self.hidden_units))

        for timestep in tf.range(self.timesteps):
            # Get the the timestep of the inputs
            x_t = tf.gather(inputs, timestep, axis=1)  # Same as x_t = inputs[:, timestep, :]

            # Compute outputs and hidden states
            y_t, h_t = self.gru_cell(x_t, h_t)
            
            # Update y_t at the t^th timestep
            self.rnn_outputs = self.rnn_outputs[:, timestep, :].assign(y_t)

        # Outputs need to have same last dimension as inputs
        outputs = self.dense(self.rnn_outputs)

        return outputs

一个会抛出错误的例子：

# Arbitrary values for dataset
num_samples = 128
batch_size = 4
timesteps = 5
features = 10

# Arbitrary dataset
x = tf.random.uniform(shape=(num_samples, timesteps, features))
y = tf.random.uniform(shape=(num_samples, timesteps, features))

train_data = tf.data.Dataset.from_tensor_slices((x, y))
train_data = train_data.shuffle(batch_size).batch(batch_size, drop_remainder=True)

# Model with arbitrary hidden units
model = CustomGruRNN(batch_size, timesteps, hidden_units=5)
model.compile(loss=tf.keras.losses.MeanSquaredError(), optimizer=tf.keras.optimizers.Adam())

急切奔跑时：

model.fit(train_data, epochs=2, run_eagerly=True)

纪元 1/2 警告：tensorflow：变量不存在梯度 ['stack_overflow_gru_rnn/gru_cell/kernel:0', 'stack_overflow_gru_rnn/gru_cell/recurrent_kernel:0', 'stack_overflow_gru_rnn/gru_cell/bias:0'] 最小化损失时。 ValueError: 未找到子字符串 ValueError

不急于奔跑时：

model.fit(train_data, epochs=2, run_eagerly=False)

纪元 1/2 类型错误：在用户代码中： TypeError：无法将 NoneType 转换为张量或操作。

【参考方案1】：

编辑：

虽然 TensorFlow 指南的答案就足够了，但我认为我的自我回答问题涉及 RNN 的自定义单元格是一个更好的选择。请参阅this answer。使用自定义 RNN 单元消除了使用 tf.Transpose 和 tf.TensorArray 的需要，从而降低了代码的复杂性，同时提高了可读性。

原始自我回答：

TensorFlow 的Guide to Effective TensorFlow2 底部附近描述的 DynamicRNN 的使用解决了我的问题。

为了简要扩展 DynamicRNN 的概念用途，定义了一个 RNN 单元，在我的例子中是 GRU，然后可以在 tf.range 循环中定义任意数量的自定义步骤。变量应该在循环外部但在调用方法本身内部使用tf.TensorArray 对象进行跟踪，并且可以通过简单地调用（输入）张量的.shape 方法来确定此类数组的大小。值得注意的是，DynamicRNN 对象在模型拟合中工作，其中默认执行模式是“Graph”模式，而不是较慢的“Eager Execution”模式。

最后，可能需要使用“DynamicRNN”，因为默认情况下，“tf.keras.layers.GRU”计算由以下循环逻辑松散地描述（假设“f”定义了一个 GRU 单元）：

# Numpy is used here for ease of indexing, but in general you should use
# tensors and transpose them accordingly (see the previously linked guide)
inputs = np.random.randn((batch, total_timesteps, features))

# List for tracking outputs -- just for simple demonstration... again please see the guide for more details
outputs = []

# Initialize the 'hidden state' (often referred to as h_naught and denoted h_0) of the RNN cell
state_at_t_minus_1 = tf.zeros(shape=(batch, hidden_cell_units))

# Iterate through the input until all timesteps in the sequence have been 'seen' by the GRU cell function 'f'
for timestep_t in total_timesteps:
    # This is of shape (batch, features)
    input_at_t = inputs[:, timestep_t, :]

    # output_at_t of shape (batch, hidden_units_of_cell) and state_at_t (batch, hidden_units_of_cell)
    output_at_t, state_at_t = f(input_at_t, state_at_t_minus_1)
    outputs.append(output_at_t)

    # When the loop restarts, this variable will be used in the next GRU Cell function call 'f'
    state_at_t_minus_1 = state_at_t

可能希望在循环逻辑的 for 循环中添加其他步骤（例如，密集层、其他层等）以修改传递给 GRU 单元函数“f”的输入和状态。这是 DynamicRNN 的动机之一。