TVMとKerasで試す量子化ニューラルネットワーク
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
Author And Source
この問題について(TVMとKerasで試す量子化ニューラルネットワーク), 我々は、より多くの情報をここで見つけました https://qiita.com/yoku_/items/03d3639cc0ebe29a88e7著者帰属:元の著者の情報は、元のURLに含まれています。著作権は原作者に属する。
Content is automatically searched and collected through network algorithms . If there is a violation . Please contact us . We will adjust (correct author information ,or delete content ) as soon as possible .