Как сделать свертку с fp16(Eigen::half) на тензорном потоке

Question

Как я могу использовать tenorflow, чтобы сделать свертку, используя fp16 на GPU? (Python API, используя __half или Eigen::half).

Я хочу протестировать модель с fp16 на тензорном потоке, но я застрял. На самом деле я обнаружил, что свертка fp16 в тензорном потоке выглядит как приведение результата свертки fp32 в fp16, а это не то, что мне нужно.

Я попытался дать tf.nn.conv2d вход fp16 в формате fp16, и дать tf.nn.conv2d вход fp16 в формате fp32 (tf.cast его в 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 не должен быть одинаковым (на мой взгляд). Как я могу использовать tenorflow, чтобы сделать свертку, используя реальный fp16 на GPU? (API Python, использующий __half или Eigen:: half)

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 регистрирует ядра для fp16 для CPU и для GPU, поэтому нет оснований думать, что делает что-то еще. У меня нет большого опыта работы с fp16, поэтому я не уверен, является ли нулевая разница "нормальной", но, похоже, нет никакого способа, которым conv16 использует что-либо, кроме свертки fp16.

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.