Tensorflow LSTM/GRU 每个时期重置状态一次，而不是每个新批次

Question

我基于 GRU 训练以下模型，注意我将参数 stateful=True 传递给 GRU 构建器。

class LearningToSurpriseModel(tf.keras.Model):
  def __init__(self, vocab_size, embedding_dim, rnn_units):
    super().__init__(self)
    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
    self.gru = tf.keras.layers.GRU(rnn_units,
                                   stateful=True,
                                   return_sequences=True,
                                   return_state=True,
                                   reset_after=True  
                                   )
    self.dense = tf.keras.layers.Dense(vocab_size)

  def call(self, inputs, states=None, return_state=False, training=False):
    x = inputs
    x = self.embedding(x, training=training)
    if states is None:
      states = self.gru.get_initial_state(x)
    x, states = self.gru(x, initial_state=states, training=training)
    x = self.dense(x, training=training)

    if return_state:
      return x, states
    else:
      return x

  @tf.function
  def train_step(self, inputs):
    [defining here my training step]

我实例化我的模型

model = LearningToSurpriseModel(
    vocab_size=len(ids_from_chars.get_vocabulary()),
    embedding_dim=embedding_dim,
    rnn_units=rnn_units
    )

[编译和做事]

训练EPOCHS epochs

for i in range(EPOCHS):
  model.fit(train_dataset, validation_data=validation_dataset, epochs=1, callbacks = [EarlyS], verbose=1)
  model.reset_states()

这段代码关于 GRU 状态的行为是什么：是为每批新数据更新状态还是只为每个新纪元更新状态？所需的行为是仅对每个新纪元进行重置。如果没有完成，如何实现？

编辑

Tensorflow 为Models 实现了reset_states 函数

  def reset_states(self):
    for layer in self.layers:
      if hasattr(layer, 'reset_states') and getattr(layer, 'stateful', False):
        layer.reset_states()

这是否意味着（与其他文档似乎暗示的相反）只有在 stateful=False 时才能重置状态？这是我从getattr(layer, 'stateful', False) 的条件推断出来的。

Answer 1

您可以尝试在自定义Callback 中重置状态：

model = LearningToSurpriseModel(
    vocab_size=len(ids_from_chars.get_vocabulary()),
    embedding_dim=embedding_dim,
    rnn_units=rnn_units
    )

gru_layer = model.layers[1]

class CustomCallback(tf.keras.callbacks.Callback):
   def __init__(self, gru_layer):
        self.gru_layer = gru_layer
   def on_epoch_end(self, epoch, logs=None):
        self.gru_layer.reset_states()

model.fit(train_dataset, validation_data=validation_dataset, epochs=1, callbacks = [EarlyS, CustomCallback(gru_layer)], verbose=1)

另外，请参阅post，了解何时重置 GRU 状态。