Почему этот RNN в тензорном потоке не учится?

Question

Я пытаюсь обучить RNN без использования RNN API в тензорном потоке (2) в Python 3.7, поэтому код очень прост. Что-то идет не так, но я не уверен, что это такое.

В качестве справки я использую набор данных из этого учебника по тензорному потоку , поэтому я знаю, что ошибка должна приблизительно сходитьсяк. Мой код RNN следующий. Он пытается использовать предыдущие 20 временных шагов, чтобы предсказать значение ряда на 21-м временном шаге. Я тренируюсь партиями размером 256.

Хотя потери со временем уменьшаются, потолок примерно в 10 раз больше, чем если бы я следовал учебному подходу. Может ли быть какая-то проблема с обратным распространением во времени?

state_size = 20  #dimensionality of the network
BATCH_SIZE = 256
#define recurrent weights and biases. W has 1 more dimension that the state 
#dimension as also processes the inputs
W = tf.Variable(np.random.rand(state_size+1, state_size), dtype=tf.float32)
b = tf.Variable(np.zeros((1,state_size)), dtype=tf.float32)

#weights and biases for the output
W2 = tf.Variable(np.random.rand(state_size, 1),dtype=tf.float32)
b2 = tf.Variable(np.zeros((1,1)), dtype=tf.float32)


init_state = tf.Variable(np.random.normal(size=[BATCH_SIZE,state_size]),dtype='float32')
optimizer = tf.keras.optimizers.Adam(1e-3)
losses = []
for epoch in range(20):
    with tf.GradientTape() as tape:
        loss = 0

        for batch_idx in range(200):

            current_state = init_state

            batchx = x_train_uni[batch_idx*BATCH_SIZE:(batch_idx+1)*BATCH_SIZE].swapaxes(0,1)
            batchy = y_train_uni[batch_idx*BATCH_SIZE:(batch_idx+1)*BATCH_SIZE]
            #forward pass through the timesteps
            for x in batchx:
                inst = tf.concat([current_state,x],1)  #concatenate state and inputs for that timepoint
                current_state = tf.tanh(tf.matmul(inst, W) + b)  #
            #predict using the hidden state after the full forward pass
            pred = tf.matmul(current_state,W2) + b2

            loss += tf.reduce_mean(tf.abs(batchy-pred))

    #get gradients with respect to parameters
    gradients = tape.gradient(loss, [W,b,W2,b2])
    #apply gradients
    optimizer.apply_gradients(zip(gradients, [W,b,W2,b2]))
    losses.append(loss)
    print(loss)