tensorflow - 寻找跨多个 GPU 训练的 tf slim 代码示例

Question

我正在寻找实现跨多个 GPU 训练的 tf slim 代码示例。

具体来说：1）如何在多个 GPU 之间分配工作负载？2）在变量共享/范围方面我需要注意什么吗？3）我需要做些什么来从检查点恢复变量吗？

我在 Stack Overflow 和 github 上搜索了示例，但似乎找不到任何示例（只是更多问题）。但是，如果这已经在某个地方得到了回答，我深表歉意 - 如果您指出我正确的方向，我将不胜感激。

谢谢大家！

Answer 1

tf-slim 已经包含代码/示例，尽管它到处都是并且不容易理解。我将在下面概述微调模型的基本步骤：

1. 按照此处自述文件中的说明设置 slim
2. 下载像 inception_v3 和鲜花数据集这样的预训练模型（slim 的数据集文件夹中的脚本）

3. 运行 train_image_classifier.py 脚本：

   python train_image_classifier.py --dataset_dir=${DATASET_DIR} --train_dir=${TRAIN_DIR} --dataset_name=flowers --dataset_split=train --model_name=inception_v3 --checkpoint_path=${CHECKPOINT_PATH} --checkpoint_exclude_scopes=InceptionV3/Logits ,InceptionV3/AuxLogits --trainable_scopes=InceptionV3/Logits,InceptionV3/AuxLogits --num_clones=2

这里num_clones参数对应于单个主机上的 GPU 数量。跨多个主机运行训练有点复杂。我已经运行了这个，并且在训练时可以看到几乎 2 倍的加速。