Многослойный градиент персептрона исчезает

Question

optimizer = optimizers.Adam(lr=1e-1)

# this is weight and bias
w1 = tf.Variable(tf.random.truncated_normal([13, 256], stddev=0.1))
b1 = tf.Variable(tf.random.truncated_normal([256]))
w2 = tf.Variable(tf.random.truncated_normal([256, 128], stddev=0.1))
b2 = tf.Variable(tf.random.truncated_normal([128]))
w3 = tf.Variable(tf.random.truncated_normal([128, 1], stddev=0.1))
b3 = tf.Variable(tf.random.truncated_normal([1]))

train_db = tf.data.Dataset.from_tensor_slices((x_train, y_train))
test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))

train_db = train_db.shuffle(10000).batch(bacth_size)
test_db = test_db.batch(bacth_size)

for epochs in range(500):
    for step, (x, y) in enumerate(train_db):
        x = tf.cast(x, dtype=tf.float32)
        y = tf.cast(y, dtype=tf.float32)
        with tf.GradientTape() as tape:
            tape.watch([w1, b1, w2, b2, w3, b3])
            h1 = tf.nn.relu(x @ w1 + b1)
            h2 = tf.nn.relu(h1 @ w2 + b2)
            logit = h2 @ w3 + b3
            loss = tf.reduce_mean(tf.losses.MSE(y, logit))

        grads = tape.gradient(loss, [w1, b1, w2, b2, w3, b3])
        optimizer.apply_gradients(zip(grads, [w1, b1, w2, b2, w3, b3]))

        if step % 100 == 0:
            print(epochs, step, 'loss, ', float(loss))

h1 = tf.nn.relu(x_test@w1 + b1)
h2 = tf.nn.relu(h1@w2 + b2)
y_predict = h2 @ w3 + b3

**

Я пытаюсь реализовать многослойный персептрон, используя базовый c API Тензор потока 2.0. Это простой многослойный персептрон, но я не знаю, почему исчезает градиент. Я новичок ie в глубоком обучении, пожалуйста, помогите мне. **