TensorFlow と bfloat16
blog
2019-4-15 2:52 JST

TensorFlow が用意している数字を表す型に bfloat16 があります。bfloat16 は sign bittが1、exponent の 部分が 8, fraction が 7と非常に特殊な形をしています。

次の図は PI( π )を bfloat16 で表した図です。ご存じのとおり PI は 3.14159... ですが bfloat16 ではそこまで精度がだせず3.140625 までしか近似できません。通常の 32bit float なら 23bit 分の fraction があるのに対しfraction がわずか 7bit ですから、整数で言えば 0 ～ 127 までしか表せない訳で、その程度しか精度が保てないわけです。

ちょっとした計算も誤差が多くなるレベルと言っても過言ではありません。

TensorFlow で bfloat16

ただし、どうやら NN の世界ではこの精度で十分らしく、TensorFlow で積極的にサポートされています。NumPy でもサポートされている？らしく、NumPy 形式の bfloat16 へも TensorFlow から cast 出来ます。すでに簡単なモデルで MNIST を使って学習したデータがあるのでその値を bfloat16 に変換してみましょう。16bit なので FPGA にモデルを構築した際に利用できればメモリの節約になります。

次に既にある浮動小数点のデータを NumPy で読み、TensorFlow で bfloat16 に変換して、さらにそれを評価させることでNumPy スタイルの bfloat16 に変換してみましょう。

bfloat.py
import numpy as np
import tensorflow as tf

w_np = np.loadtxt('w_value.txt', delimiter=',')
b_np = np.loadtxt('b_value.txt', delimiter=',')

print(type(w_np), type(w_np[0][0]))

w_b = tf.cast(w_np, tf.bfloat16)
b_b = tf.cast(b_np, tf.bfloat16)

print(type(w_b), type(w_b[0][0]))

with tf.Session() as sess:
    w_b_np = w_b.eval()
    b_b_np = b_b.eval()

print(type(w_b_np[0][0]))
#np.savetxt('w_b_value.txt', w_b_np, delimiter=',')
#np.savetxt('b_b_value.txt', b_b_np, delimiter=',')
<class 'numpy.ndarray'> <class 'numpy.float64'>
<class 'tensorflow.python.framework.ops.Tensor'>
    <class 'tensorflow.python.framework.ops.Tensor'>
<class 'bfloat16'>

NumPy で バイナリ化

最終的には 16bit の"バイナリ"にしたいので savetxt ではなく save でセーブします。

bfloat.py
np.save('w_b_value.npy', w_b_np)
np.save('b_b_value.npy', b_b_np)

生成された npy は余計なヘッダがついているので dd でスキップして素の"バイナリ"に変換します。あとは od でテキストにすれば、python の list にまで変換可能です。最終的には w_bin.py と b_bin.py ができあがります。

NumPy で評価してみる

float32 で計算したパラメタを bfloat16 にまで落として同じモデルに再学習無しで使えるでしょうか？試しに、bfloat16 を再度 float に変換して(この時点で精度が落ちる)そのパラメタをもって NumPy で計算させてみます。自家製の bfloat2float という関数を使っています。

結果は勿論、うまくいきました。

numpy-test.py
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)

from tensorflow.python import pywrap_tensorflow

bfloat16 = pywrap_tensorflow.TF_bfloat16_type()
import numpy as np
from b_bin import B_PARAM
from w_bin import W_PARAM
from float2bfloat import float2bfloat
from float2bfloat import bfloat2float

b = np.array(list(map(lambda x: bfloat2float(x), B_PARAM)))
w = np.array(list(map(lambda x: bfloat2float(x), W_PARAM)))

w = w.reshape(784, 10)

img=mnist.test.images[0]
print(img.shape, w.shape)

xw = img @ w
print(xw.shape)
rv = xw + b

print(rv)

おまけ:bfloat2float

Python なんで、この辺の変換は簡単に書けます。この程度の量なら Python で十分通用します。

bfloat2float.py
def bfloat2float(x):
    s = 1 if x & 0x8000 else 0
    e = (x >> 7) & 0xFF
    n = x & 0x7F
    fn = n / (2 ** 7) + 1
    fn_e = fn * (2 ** (e - 127))
    return fn_e if s == 0 else -fn_e

リンク集

NumPy — NumPy
NumPy のウェブサイト。pip で入れるだけで簡単に使えるので滅多に行かない。
2019-4-14 2:39 JST
Copyright 2001-2019 Sinby Corporation.