如何使检查点在tf.train优化器中存储时刻和其他相关变量

optimizer = tf.train.AdamOptimizer(learning_rate)
model = MyModel()

checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)

checkpoint_dir = './checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint_manager = tf.train.CheckpointManager(checkpoint, checkpoint_dir, max_to_keep=5)

# 在每个训练步骤结束后保存检查点
checkpoint_manager.save()

# 恢复最新的检查点
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

# 或者恢复指定的检查点
checkpoint.restore('./checkpoints/ckpt-10')