tensorflow - Tesla V100 上的 TF1.4 未启用混合精度

Question

我有兴趣测试我的神经网络(一个用作生成器的自动编码器 + 一个用作鉴别器的 CNN)，它使用 3dconv/deconv 层和新的 Volta 架构，并从混合精度训练中受益。我用 CUDA 9 和 CudNN 7.0 编译了 Tensorflow 1.4 的最新源代码，并将我的 conv/deconv 层使用的所有可训练变量转换为 tf.float16。此外，我所有的输入和输出张量的大小都是 8 的倍数。

不幸的是，我没有看到这种配置有任何实质性的速度提升，训练时间与使用 tf.float32 时大致相似。我的理解是，使用 Volta 架构和 cuDNN 7.0，混合精度应该会被 TF 自动检测到，从而能够使用 Tensor Core 数学。我错了，还是我应该做些什么来启用它？
我也试过TF1.5 nighlty build，好像比我自定义的1.4还要慢。

如果任何参与 Tensorflow 的开发人员能回答这个问题，我将不胜感激。

编辑:在与 NVIDIA 技术支持交谈后，似乎在支持 float16 的同时，TF 为简单的 2D conv Ops 集成了混合精度加速，但目前还没有为 3D conv Ops 集成。

Answer 1

基于 NVIDIA documentation我使用 FP16 (TensorCore) 运行基准测试。为此，我修改了 alexnet_benchmark由 tensorflow 提供:
https://gist.github.com/melgor/946b9643aa25dd3839a86804fc580741

总体而言，AlexNet 仅快 35%，而不是那么多。我希望能快 2 倍。此外，也许 Resnet 会产生更大的不同。好消息是我可以用 batch_size = 5120 拟合模型(fp32 不能)，一次 FB pass 需要 0.653，所以训练 ImageNet 90 epochs 需要大约 4 小时。
batch_size=512 alexnet_fp32: Forward-backward across 100 steps, 0.099 +/- 0.000 sec / batch alexnet_fp16: Forward-backward across 100 steps, 0.064 +/- 0.000 sec / batch
编辑:

我设法在 FP16 上运行 ResNet 模型(但没有 BatchNorm，由于某种原因 BN 不适用于 fp16):
batch_size=256 resnet50_fp32: Forward-backward across 100 steps, 0.575 +/- 0.001 sec / batch resnet50_fp16: Forward-backward across 100 steps, 0.504 +/- 0.001 sec / batchbatch_size=128 resnet152_fp32: Forward-backward across 100 steps, 0.757 +/- 0.001 sec / batch resnet152_fp16: Forward-backward across 100 steps, 0.581 +/- 0.010 sec / batch
ResNet 的增益甚至更小。看起来 FP16 在 V100 上没有很多增益，不知道为什么。可能目前对 TensorCore 的支持还没有完全整合。