有没有一种异步编写TensorFlow检查点的方法？

Question

目前，我在训练期间创建检查点，如下所示(伪代码)：

while(training):
    model.train()

    if it_is_time_for_validation():
        metrics = model.validate()

        if metrics.are_good():
             saver = tf.train.Saver()
             res = saver.save(sess=session, save_path=checkpoint_file_path)

Saver.save方法阻塞I/O，阻止运行下一次迭代。我的模型的重量大小是几百兆字节，写这些东西需要一段时间。

根据我的计算，根据检查点频率，总体上，GPU花费5-10%的时间等待检查点完成，而不是进行有用的计算。(5-10%相当于一天的计算时间)

有没有办法异步执行检查点，以减少计算时间的浪费？

实现草图:首先，我们可以将所有必需的东西从设备内存复制到主机，并在单独的线程上执行磁盘I/O。Saver.save将在memcopy之后返回，而无需等待磁盘操作，因为现在训练设备副本而不会搞乱检查点是安全的。如果前一次迭代中有I/O挂起，Saver.save仍会在重新进入时阻塞。

我认为它目前还没有实现，所以我也对可能的变通方法感兴趣。这个想法是否足够好，可以作为GitHub上的一个特性请求？

Answer 1

您可以通过在单独的线程中运行saver.save()来异步编写检查点。(内部) SVTimerCheckpointThread是一个在训练后台定期运行saver.save()的代码示例。请注意，tf.train.Supervisor是一个实用程序类，可以帮助管理这样的后台线程(也用于编写TensorBoard摘要日志等)，因此您可能希望使用它。