python - 用于 VGG19 模型参数的 Tensorflow Float16

Question

我正在Python 中使用 TensorFlow（带有 19.01 Nvidia TensorFlow Docker 映像）进行神经风格迁移。我有一个 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 来减少模型内存占用？

标签： pythontensorflowkeras

Answer 1

TensorFlow 有多种管理混合精度的方法。最合适的混合精度方法取决于您计划使用的优化器。例如，Keras 优化器有一个API，旨在轻松地以一种或另一种方式移植代码。它被称为mixed_precision。

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