tensorflow - 使用 GPU 功能在 tensorflow-hub 上使用 retrain.py 重新训练图像

Question

我是 Tensorflow 的新手，我正在使用 retrain.py 来训练一些图像。如果我有一个包含 10000 张图像的更大数据库并且我有一个支持 GPU 的系统。如何使用 retrain.py 在我的 Nvidia GPU 上运行。这样训练会更快完成。

我正在按照以下链接中的步骤操作

https://www.tensorflow.org/hub/tutorials/image_retraining

Answer 1

要获得 GPU 支持，请务必安装 PIP 包 tensorflow-gpu 而不是普通的 tensorflow。您应该会从中看到 retrain.py 的一些性能优势。也就是说，retrain.py 显示了它的年龄（远早于 TF Hub）并且没有很好地利用 GPU，因为它在提取瓶颈值时没有正确地批处理图像。

如果您准备好生活在 TF 2.0.0alpha0（上周 宣布）的前沿，请查看 Hub 的示例/colab/tf2_image_retraining.ipynb，它更小、更快（如果您使用 GPU），甚至支持微调图像模块。