Tensorflow：是否可以修改检查点中的全局步骤

Question

我正在尝试修改检查点中的global_step，以便将训练从一台机器转移到另一台机器。

假设我在机器 A 上训练了几天。现在我买了一台新机器 B，它有更好的显卡和更多的 GPU 内存，想把训练从机器 A 转移到机器 B。

为了恢复机器B中的检查点，我之前在机器A的Saver.save中指定了global_step，但batch_size更小，sub_iterations更大。

batch_size=10
sub_iterations=500
for (...)
    for i in range(sub_iterations):
        inputs, labels = next_batch(batch_size)
        session.run(optimizer, feed_dict={inputs: inputs, labels: labels})

saver = tf.train.Saver()
saver.save(session, checkpoints_path, global_step)

现在我将包括检查点在内的所有文件从机器 A 复制到机器 B。由于机器 B 有更多 GPU 内存，我可以将 batch_size 修改为更大的值，但使用更少的 sub_iterations。

batch_size=100
sub_iterations=50 # = 500 / (100/10)
for (...)
    for i in range(sub_iterations):
        inputs, labels = next_batch(batch_size)
        session.run(optimizer, feed_dict={inputs: inputs, labels: labels})

但是我们不能直接恢复复制的检查点，因为global_step在机器B中是不同的。例如，tf.train.exponential_decay将产生不正确的learning_rate，因为机器B中sub_iterations的数量减少了。

learning_rate = tf.train.exponential_decay(..., global_step, sub_iterations, decay_rate, staircase=True)

是否可以修改检查点中的global_step？还是有其他更合适的方法来处理这种情况？

编辑 1

除了计算learning_rate之外，我还使用了global_step来计算和减少迭代次数。

while i < iterations:
    j = 0

    while j < sub_iterations:
        inputs, labels = next_batch(batch_size)
        feed_dict_train = {inputs: inputs, labels: labels}
        _, gstep = session.run([optimizer, global_step], feed_dict=feed_dict_train)
        if (i == 0) and (j == 0):
            i, j = int(gstep/ sub_iterations), numpy.mod(gstep, sub_iterations)
        j = j + 1
    i = i + 1

我们将从新的i 和j 开始迭代。请随时对此发表评论，因为这可能不是恢复检查点并从加载的检查点继续训练的好方法。

编辑 2

在机器 A 中，假设 iterations 是 10,000，sub_iterations 是 500，batch_size 是 10。所以我们要训练的批次总数是 10000x500x10 = 50,000,000。假设我们已经训练了几天，global_step 变成了501。所以训练的batch总数为501x10 = 5010。剩余未训练的batch数量为5011到50,000,000。如果我们应用i, j = int(gstep/ sub_iterations), numpy.mod(gstep, sub_iterations)，i 的最后训练值是 501/500=1，j 是 501%500=1。

现在您已将包括检查点在内的所有文件复制到机器 B。由于 B 具有更多 GPU 内存，我们可以为一个子迭代训练更多批次，我们将 batch_size 设置为 100，并将 sub_iterations 调整为50 但将iterations 保留为 10000。要训练的总批次数仍然是 50,000,000。那么问题来了，我们如何才能开始训练从 5011 到 50,000,000 的批次，并且不对前 5010 个样本再次进行训练？

为了在机器 B 中从 5011 开始和训练批次，我们应该将 i 设置为 1 并将 j 设置为 0，因为它已训练的总批次将是 (1*50+0)*1​​00 = 5,000 接近 5,010（因为 batch_size 在机器 B 中是 100 而在机器 A 中是 10，所以我们不能完全从 5,010 开始，我们可以选择 5,000 或 5,100）。

如果我们不调整global_step（如@coder3101所建议的那样），并在机器B中使用回i, j = int(gstep/ sub_iterations), numpy.mod(gstep, sub_iterations)，i将变为501/50=10，j将变为501%50= 1.所以我们将从不正确的批次 50,100 (=501*batch_size=501*100) 开始训练。

之所以引入这个公式i, j = int(gstep/ sub_iterations), numpy.mod(gstep, sub_iterations)是因为如果我们停止机器A的训练，我们可以恢复检查点并使用这个公式继续机器A的训练。但是，当我们将训练从机器 A 转移到机器 B 时，似乎这个公式不适用。因此我希望修改检查点中的 global_step 来处理这种情况，并想知道这是否可能。

Answer 1

是的。这是可能的。

要修改机器B中的global_step，您必须执行以下步骤：

计算对应的global_step

在上面的例子中，机器A中的global_step是501，训练的batch总数是501x10=5010。所以机器B中对应的global_step是5010/100=50。

修改检查点文件名

修改 model_checkpoint_path 和 all_model_checkpoint_paths 值的后缀以在 /checkpoint 中使用正确的步骤编号（例如上例中的 50）。

修改检查点的索引、元数据和数据的文件名以使用正确的步骤号。

恢复检查点后覆盖global_step

saver.restore(session, model_checkpoint_path)
initial_global_step = global_step.assign(50)
session.run(initial_global_step)
...
# do the training

现在如果你恢复检查点（包括global_step）并覆盖global_step，训练将使用更新的global_step并正确调整i和j和learning_rate。

Answer 2

如果您希望机器 B 上的衰减学习率与 A 上的相同，您可以调整函数 tf.train.exponential_dacay 的其他参数。例如，您可以在新机器上更改衰减率。

为了找到decay_rate，您需要知道exponential_decay 是如何计算的。

decay_learning_rate = learning_rate * decay_rate ^ (global_step/decay_steps)

global_step/decay_step 如果 staircase==True 则产生一个纯整数

decay_steps 是您的 sub_iteration，global_step 来自机器 A。

您可以更改 decay_rate，使机器 B 上的新学习率与您对机器 A 在 global_step 的预期相同或接近。

另外，您可以更改机器 B 的初始学习率，以实现从机器 A 到 B 的均匀指数速率衰减。

简而言之，从 A 导入到 B 时，您更改了 1 个变量 (sub_iteration)，并保持 global_step 不变。您可以调整 exponential_decay(..) 的其他 2 个变量，使函数的输出学习率与机器 A 在 global_step 的预期相同。