TVMとKerasで試す量子化ニューラルネットワーク #Python

TVM v0.5で8bit整数型へのpost-training量子化がサポートされました。
現時点(2019-3-27)ではAPIドキュメント・チュートリアルが公開されていないため、簡単に使い方をまとめておきます。

ソースコードはこちら

手順

Keras公式のResNet50モデルを例に進めていきます。

import keras
model = keras.applications.resnet50.ResNet50(include_top=True, weights='imagenet', input_tensor=None, input_shape=None, pooling=None, classes=1000)

モデル変換

まずはこのモデルをRelayIRの関数に変換しましょう。v0.5で正式サポートされたRelayIRですが、既にKerasを含め多くのDLフレームワークからのモデル変換機能が用意されています。

import tvm.relay as relay

input_name = 'input_1'
image_shape = (1, 3 , 224, 224)  # NCWH
func, params = relay.frontend.from_keras(model, image_shape)

量子化

早速この関数を量子化してみましょう。

import tvm
import tvm.relay as relay

with relay.quantize.qconfig(global_scale=8.0):
    func_quant = relay.quantize.quantize(func, params)
    print(str(func_quant))

  ...
  %9 = round(%8)
  %10 = clip(%9, a_min=-127, a_max=127)
  %11 = cast(%10, dtype="int8")
  %12 = nn.conv2d(%11, meta[relay.Constant][4], channels=64, kernel_size=[1, 1], out_dtype="int32")
  %13 = add(%12, meta[relay.Constant][5])
  %14 = add(%13, meta[relay.Constant][6])
  %15 = nn.relu(%14)
  %16 = add(%15, 64)
  %17 = right_shift(%16, 7)
  %18 = clip(%17, a_min=-127, a_max=127)
  %19 = cast(%18, dtype="int8")
  ...

conv2dへの入力値が8bit整数化されていることが分かります。量子化アルゴリズムの詳細はこの機能のコミッタであるZiheng JiangさんのLT資料がシンプルかつ直観的にまとまっています。

コンパイル

生成されたRelay関数をコンパイルします。ターゲットはLLVM(x86 CPU)としましょう。比較のため量子化前の関数も併せてコンパイルしておきます。

target = 'llvm'

# 量子化有無による差異のみチェックするため,グラフ最適化は無効に
with relay.build_config(opt_level=0):
    modules = relay.build_module.build(func, target, params=params)
    modules_quant = relay.build_module.build(func_quant, target, params=params)

推論

一般にpost-training量子化を施したモデルは、元のモデルに比べ多少精度が低下する傾向にあります。TVMの量子化はどうでしょうか。
おなじみImageNetは tabby cat の画像でテストしてみます。

def predict_tvm(modules, data):
    # create module
    graph, lib, params = modules
    module = tvm.contrib.graph_runtime.create(graph, lib, tvm.cpu())

    # set input and parameters
    module.set_input(input_name, tvm.nd.array(data.astype(np.float32)))
    module.set_input(**params)

    # get output
    module.run()
    return module.get_output(0, tvm.nd.empty((1, 1000))).asnumpy()

def show_top5_accuracy(y_pred, synset):
    top5_ids = y_pred.flatten().argsort()[::-1][:5]

    for i, c in enumerate(top5_ids):
        print(f'{i+1} : {c}  {synset[c]}')

y_pred_tvm = predict_tvm(modules, data)
y_pred_tvm_quant = predict_tvm(modules_quant, data)

print('\nw/o quantization')
show_top5_accuracy(y_pred_tvm, synset)

print('\nwith quantization')
show_top5_accuracy(y_pred_tvm_quant, synset)

w/o quantization
1 : 278  kit fox, Vulpes macrotis
2 : 277  red fox, Vulpes vulpes
3 : 282  tiger cat
4 : 285  Egyptian cat
5 : 356  weasel

with quantization
1 : 263  Pembroke, Pembroke Welsh corgi
2 : 282  tiger cat
3 : 285  Egyptian cat
4 : 281  tabby, tabby cat
5 : 722  ping-pong ball

量子化の影響か両者の推論結果が若干変わっているものの、傾向としては似たようなクラスが出力されているようです。コーギーと子狐、どちらがよりこの茶トラの子猫に近しいかは議論の別れるところですが...

その他、いくつか別クラスの画像でも試してみましたが、概ね量子化前後で同等の推論結果が得られました。

ファイルサイズ

最後に重みファイルのサイズも比較してみましょう。

# save weight
model.save_weights('weight_keras.h5')

params = modules[2]
with open("deploy_param.params", "wb") as f:
    f.write(relay.save_param_dict(params))

params_quant = modules_quant[2]
with open("deploy_param_quant.params", "wb") as f:
    f.write(relay.save_param_dict(params_quant))

Keras(weights)	TVM(未量子化)	TVM(量子化)
99MB	98MB	31MB

ばっちり軽量化できました！実質コードに2行追加しただけで、重みのサイズが約1/3になったのは中々インパクトがあるのではないでしょうか。

推論速度

最後に推論速度を計っておきましょう。注意点として、今回のモデルはTVMのキモであるグラフ最適化を行っておらず、ターゲットデバイス向けのチューニングも施していません。絶対的な推論速度は参考程度にとどめてください。

print('w/o quantization')
%timeit -n10 module.run()

print('\nwith quantization')
%timeit -n10 module_quant.run()

w/o quantization
141 ms ± 447 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

with quantization
192 ms ± 77.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

量子化後の方が僅かに遅くなっていることが分かります。今回ターゲットとしたx86-64 CPUのようにFPU性能やキャッシュ容量が十分大きいデバイスでは、サイズの縮小による高速化よりも、量子化時のキャストやクリッピング処理による負荷の方が大きいようです。ARM CPUやFPGAでも改めて検証してみたいところですね。

終わりに

今回はTVM v0.5の新機能、post-training量子化を試してみました。既存コードに数行追加するだけで量子化にトライできる手軽さには驚かされるばかりです。既にエッジ向け推論フレームワークとして関心が高まりつつあるTVMですが、量子化の実装で今後ますます注目を集めそうです。

環境

Ubuntu 16.04 LTS
AMD Ryzen 7 1700
TVM master a7c90ee579e0e830e349f372f3783a0bd31ac605 (0.6.dev)
Python 3.6.8
- Keras 2.2.4
- tensorflow-gpu 1.11.0

参考資料

TVM Conference 2018, Automatic Quantization for TVM
https://sampl.cs.washington.edu/tvmconf/slides/11-Ziheng-Jiang.pdf

TVM tutorials, Compile Keras Models
https://docs.tvm.ai/tutorials/frontend/from_keras.html#compile-keras-models