如何在张量流上使用 fp16(Eigen::half) 进行卷积

Question

如何使用 tensorflow 在 GPU 上使用 fp16 进行卷积？(使用 __half 或 Eigen：：half 的 python api(。

我想在张量流上用 fp16 测试一个模型，但我卡住了。实际上，我发现 tensorflow 中的 fp16 卷积似乎将 fp32 卷积的结果转换为 fp16，这不是我所需要的。

我尝试给 tf.nn.conv2d 一个 fp16 格式的 fp16 输入，给 tf.nn.conv2d 一个 fp32 格式的 fp16 输入(tf.cast it 变成 fp32(，然后将结果 tf.cast 转换为 fp16，他们给出了完全相同的结果。 但正如我认为的那样，在 fp16 中进行卷积与在 fp32 中进行卷积然后将其转换为 fp16 是不同的，我错了吗？ 请帮助我，谢谢。

environment:
ubuntu 16.04
tensorflow 1.9.0
cuda 9.0
Tesla V100

import tensorflow as tf
import numpy as np
import os
def conv16_32(input, kernel): # fake fp16 convolution
input = tf.cast(input, tf.float16)
kernel = tf.cast(kernel, tf.float16)
input = tf.cast(input, tf.float32)
kernel = tf.cast(kernel, tf.float32)
out = tf.nn.conv2d(input, kernel, [1,1,1,1], padding='VALID')
out = tf.cast(out, tf.float16)
out = tf.cast(out, tf.float64)
return out
def conv16(input, kernel): # real fp16 convolution
input = tf.cast(input, tf.float16)
kernel = tf.cast(kernel, tf.float16)
out = tf.nn.conv2d(input, kernel, [1,1,1,1], padding='VALID')
out = tf.cast(out, tf.float64)
return out
x = np.random.rand(16, 32, 32, 16).astype('float64')
w = np.random.rand(3, 3, 16, 16).astype('float64')
x = tf.get_variable('input', dtype=tf.float64, initializer=x)
w = tf.get_variable('weight', dtype=tf.float64, initializer=w)
out_16 = conv16(x, w)
out_16_32 = conv16_32(x, w)
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config = config)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
print(sess.run(tf.reduce_max(out_16_32 - out_16)))

上述两个函数给出相同的结果，假设最终的"打印"结果为零。

fp16 卷积和 fp32 卷积的结果不应该相同(在我看来(。 如何使用张量流在 GPU 上使用真正的 fp16 进行卷积？(使用 __half 或 Eigen：：half 的 Python API(

Answer 1

我认为您正确地使用了这些操作。在您的示例中，您可以检查卷积操作是否确实具有正确的类型。

conv2d_op_16 = out_16.op.inputs[0].op
print(conv2d_op_16.name, conv2d_op_16.type, conv2d_op_16.get_attr('T'))
# Conv2D Conv2D <dtype: 'float16'>
conv2d_op_16_32 = out_16_32.op.inputs[0].op.inputs[0].op
print(conv2d_op_16_32.name, conv2d_op_16_32.type, conv2d_op_16_32.get_attr('T'))
# Conv2D_1 Conv2D <dtype: 'float32'>

TensorFlow确实为CPU和GPU的fp16注册了内核，因此没有理由认为正在做其他事情。我对 fp16 没有太多经验，所以我不确定零差异是否"正常"，但除了 fp16 卷积之外，conv16似乎没有任何方法可以使用任何其他方法。

Answer 2

我试图弄清楚同样的事情。下面是一些简单的代码，您可以使用它们来测试卷积：

import tensorflow as tf
tf.enable_eager_execution()
input = tf.cast([[[[65519], [65519], [65519], [65519]]]], tf.float16) #BHWC
filter = tf.cast([[[[65519]], [[-65519]]]], tf.float16) #HWIO
tf.print(tf.nn.conv2d(input, filter, [1,1,1,1], "VALID"))

如果卷积是在 fp16 中完成的，这应该会溢出，但实际上并没有在 Tensorflow 中溢出。我得到的结果是[[[[0][0][0]]]]，这表明卷积是在 fp32 中执行的。

编辑：解决方案是设置环境变量：

TF_FP16_CONV_USE_FP32_COMPUTE=0

这给出了[[[[inf][inf][inf]]]]的结果，表明这次卷积是在fp16中进行的。似乎您至少需要一个 10x0 GPU。