用于 VGG19 模型参数的 Tensorflow Float16

Question

我正在使用 Python 中的 TensorFlow（带有 19.01 Nvidia TensorFlow Docker 映像）做 neural style transfer。我有一个 Nvidia 2070 显卡，当我尝试运行更新版本的 TensorFlow docker 映像（例如 19.08）时，出现内存不足错误。所以我决定也许是时候考虑使用 16 位精度而不是 32 位来存储 VGG19 CNN 的参数了。

我在构建机器时的初步研究使我相信从 32 切换到 16 是轻而易举的事，但现在我正在积极尝试进行转换，这不是我的经验。

这就是我所做的：

1. 我设置了 tf.keras.backend.set_floatx('float16')
2. 设置 tf.keras.backend.set_epsilon(1e-4)
3. 将 VGG19 网络的图像输入更改为 float16，并将我的代码中使用 float32 数据类型与 float16 结合使用的任何其他杂项部分。

当我运行代码时，nvidia-smi 仍然报告说我的 GPU 基本上 100% 都在使用。有没有人通过在 TensorFlow 中切换到 float16 来减少模型内存占用？

Answer 1

TensorFlow 有多种管理混合精度的方法。最合适的混合精度方法取决于您计划使用的优化器。例如，Keras 优化器有一个API designed to easily port code one way or the other。它被称为mixed_precision。

如果您使用的是 TFv1 优化器，或者 TensorFlow 提供的其他非 keras 优化器之一，您可以使用他们的graph rewrite function 将图形的各个部分转换为浮点 16。