train_step = tf.train.AdamOptimizer(1e-6).minimize(cost)改写问TensorFlow2代码

在 TensorFlow 2.x 中，您可以使用 `tf.keras.optimizers.Adam()` 来替换 `tf.train.AdamOptimizer()`。

例如，如果您想要使用 Adam 优化器来最小化成本函数 `cost`，可以使用以下代码：

import tensorflow as tf

# 创建一个 Adam 优化器，并使用 1e-6 的学习率
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-6)

# 定义一个计算梯度和更新权重的函数
@tf.function
def train_step(inputs, targets):
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss = cost(predictions, targets)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

上述代码将创建一个 Adam 优化器，并使用 `train_step()` 函数来计算梯度并更新权重。请注意，我们使用 `tf.function` 将函数转换为 TensorFlow 计算图，以提高性能。最后，我们返回损失以跟踪模型的训练进度。