tensorflow - 如何加载 TensorFlow 中的最后一个检查点？

Question

我正在 this tutorial 上练习 TensorFlow . evaluate 函数依赖于加载最新检查点的训练:

checkpoint_path = "./checkpoints/train"
ckpt = tf.train.Checkpoint(encoder=encoder,
                           decoder=decoder,
                           optimizer = optimizer)
ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)

start_epoch = 0
if ckpt_manager.latest_checkpoint:
  start_epoch = int(ckpt_manager.latest_checkpoint.split('-')[-1])
  ckpt.restore(ckpt_manager.latest_checkpoint)

for epoch in range(start_epoch, EPOCHS):
    start = time.time()
    total_loss = 0

    for (batch, (img_tensor, target)) in enumerate(dataset):
        batch_loss, t_loss = train_step(img_tensor, target)
        total_loss += t_loss

        if batch % 100 == 0:
            print ('Epoch {} Batch {} Loss {:.4f}'.format(
              epoch + 1, batch, batch_loss.numpy() / int(target.shape[1])))
    loss_plot.append(total_loss / num_steps)

    ckpt_manager.save()

如果没有 ckpt_manager.save()，evaluation 函数将不起作用。

当我们已经训练了一个模型并且检查点在 checkpoint_path 中可用时。我们应该如何在没有训练的情况下加载模型？

Answer 1

您可以使用 tf.train.latest_checkpoint 获取最新的检查点文件，然后使用 ckpt.restore 手动加载它:

checkpoint_path = "./checkpoints/train"
ckpt = tf.train.Checkpoint(encoder=encoder,
                           decoder=decoder,

ckpt_path = tf.train.latest_checkpoint(checkpoint_path)
ckpt.restore(ckpt_path)