使用 Tensorflow 构建 GRUcell 以获取隐藏状态

如何解决使用 Tensorflow 构建 GRUcell 以获取隐藏状态

我想构建一个带有一个隐藏层的 GRU 单元，但我想在每个时间步获取隐藏状态。我想训练 GRU 单元 10 次，并且每一步都得到隐藏状态。所以我尝试使用 Tensorflow 中的 GRUcell。

Target_tensor 是一个带有目标变量的张量。假设我还定义了损失函数，即代码中的 squareLoss。 hiddenStateSize 是隐藏单元的数量。

gru = tf.keras.layers.GRUCell(
    units = hiddenStateSize,activation = 'sigmoid'
)

out = tf.keras.layers.Dense(
    units = 1,activation = None
)


optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)

def train(state,X):

hidden_state,_ = gru(np.expand_dims(X,0),state)
pred = tf.squeeze(out(hidden_state)) 

return pred,hidden_state



with tf.GradientTape(persistent=True) as tape:


state = gru.get_initial_state(
        batch_size = 1,dtype = tf.float32
    )
          
for i in range(10):     
    
    pred,state = train(state,X[i])
    
    yTrue = target_tensor[i]
    loss = squareLoss(pred,yTrue)
    
    trainableVars = gru.trainable_variables + out.trainable_variables
            
    grads = tape.gradient(loss,trainableVars)

    optimizer.apply_gradients(zip(grads,trainableVars))

问题：

1. 使用 gru 和 out 是否正确？
2. 我是否正确地执行了 GradientTape？因为随着迭代的进行，损失函数的值会越来越大。