[TVMチュートリアル] ARM CPU向け畳み込みネットワークの自動チューニング

Apache TVMは、CPU、GPU、そして様々な機械学習アクセラレーションチップに適したディープラーニングコンパイラフレームワークです。TVMの中国語版ドキュメントは→ https://tvm.hyper.ai/ をご覧ください。

著者: Lianmin Zheng、Zhao Wu、Eddie Yan

特定のARMデバイス向けの自動チューニングは、最適なパフォーマンスを実現するために不可欠です。この記事では、畳み込みネットワーク全体をチューニングする方法について説明します。

TVMでは、ARM CPU向けの演算子実装はテンプレート形式で記述されており、多くの調整可能なパラメータ（タイル係数、ベクトル化、アンローリングなど）が含まれています。ニューラルネットワーク内のすべての畳み込み演算子と深度方向畳み込み演算子をチューニングした後、必要なすべての演算子の最適なパラメータ値を格納するログファイルが生成されます。TVMコンパイラはこれらの演算子をコンパイルする際に、このログファイルを照会して最適なパラメータ値を取得します。

ARMデバイス向けのプリセットパラメータもいくつか公開しました。結果はARM CPUベンチマークでご確認いただけます。

このチュートリアルはWindowsまたは最新バージョンのmacOSでは動作しませんのでご注意ください。実行するには、チュートリアルの本体を `if name == "__main__":` ブロックで囲む必要があります。

依存関係をインストールする

TVM で autotvm パッケージを使用するには、追加の依存関係をインストールする必要があります (Python 2 を使用している場合は、「3」を「2」に変更してください)。

 pip3 install --user psutil xgboost tornado cloudpickle

チューニング中のTVMパフォーマンスを向上させるため、TVMのFFIとしてCythonを使用することをお勧めします。TVMのルートディレクトリで、以下のコマンドを実行してください（Python 2を使用している場合は、「3」を「2」に置き換えてください）。

 pip3 install --user cython sudo make cython3

Python コードでパッケージをインポートします。

 import os import numpy as np import tvm from tvm import relay, autotvm import tvm.relay.testing from tvm.autotvm.tuner import XGBTuner, GATuner, RandomTuner, GridSearchTuner from tvm.contrib.utils import tempdir import tvm.contrib.graph_executor as runtime

ネットワークを定義する

まず、リレーフロントエンドAPIでネットワークを定義する必要があります。relay.testingから定義済みのネットワークを読み込むことも、MXNet、ONNX、TensorFlowからモデルを読み込むこともできます。

 def get_network(name, batch_size): """获取网络的符号定义和随机权重""" input_shape = (batch_size, 3, 224, 224) output_shape = (batch_size, 1000) if "resnet" in name: n_layer = int(name.split("-")[1]) mod, params = relay.testing.resnet.get_workload( num_layers=n_layer, batch_size=batch_size, dtype=dtype ) elif "vgg" in name: n_layer = int(name.split("-")[1]) mod, params = relay.testing.vgg.get_workload( num_layers=n_layer, batch_size=batch_size, dtype=dtype ) elif name == "mobilenet": mod, params = relay.testing.mobilenet.get_workload(batch_size=batch_size) elif name == "squeezenet_v1.1": mod, params = relay.testing.squeezenet.get_workload( batch_size=batch_size, version="1.1", dtype=dtype ) elif name == "inception_v3": input_shape = (batch_size, 3, 299, 299) mod, params = relay.testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype) elif name == "mxnet": # MXNet 模型的示例from mxnet.gluon.model_zoo.vision import get_model block = get_model("resnet18_v1", pretrained=True) mod, params = relay.frontend.from_mxnet(block, shape={"data": input_shape}, dtype=dtype) net = mod["main"] net = relay.Function( net.params, relay.nn.softmax(net.body), None, net.type_params, net.attrs ) mod = tvm.IRModule.from_expr(net) else: raise ValueError("Unsupported network: " + name) return mod, params, input_shape, output_shape

RPCトラッカーを起動する

TVMはRPCセッションを使用してARMボードと通信します。チューニング中、チューナーは生成されたコードをボードに送信し、ボード上でのコードの速度をテストします。

TVMはチューニングを高速化するために、RPCトラッカー（集中型コントローラーノード）を使用して分散デバイスを管理します。例えば、携帯電話が10台ある場合、それらをすべてトラッカーに登録し、10個のテストを並列実行することで、チューニングプロセスを高速化できます。

RPCトラッカーを起動するには、ホストマシンで次のコマンドを実行します。トラッカーはチューニングプロセス全体を通して必要なので、このコマンドを実行するには新しいターミナルを開く必要があります。

 python -m tvm.exec.rpc_tracker --host=0.0.0.0 --port=9190

期待される出力:

 INFO:RPCTracker:bind to 0.0.0.0:9190

デバイスをRPC Trackerに登録する

次のステップは、デバイスをトラッカーに登録することです。最初のステップは、ARMデバイス用のTVMランタイムをビルドすることです。

Linux の場合: デバイス上に TVM ランタイムを構築するチュートリアルに従い、デバイスをトラッカーに登録します。

 python -m tvm.exec.rpc_server --tracker=[HOST_IP]:9190 --key=rk3399

（[HOST_IP]をホストのIPアドレスに置き換えてください）

Androidの場合：以下の手順に従って、TVM RPC APKをAndroidデバイスにインストールし、Android RPCテストに合格することを確認してください。調整中は、スマートフォンの開発者向けオプションを有効にし、「変更中は画面をスリープ解除しない」にチェックを入れ、スマートフォンを電源に接続してください。

デバイスを登録した後、rpc_tracker をチェックして登録が成功したかどうかを確認します。

 python -m tvm.exec.query_rpc_tracker --host=0.0.0.0 --port=9190

たとえば、Huawei Mate 10 Pro が 2 台、Raspberry Pi 3B が 11 台、RK3399 が 2 台ある場合、出力は次のようになります...

 Queue Status ---------------------------------- key total free pending ---------------------------------- mate10pro 2 2 0 rk3399 2 2 0 rpi3b 11 11 0 ----------------------------------

複数のデバイスをトラッカーに登録すると、最適化テストが高速化されます。

最適化オプションを設定する

最適化の前に、システムを設定してください。この例ではRK3399ボードを使用しています。お使いのデバイスに合わせて`target`と`device_key`を変更してください。Androidスマートフォンを使用している場合は、`use_android`を`True`に設定してください。

 #### 设备配置#### # 将"aarch64-linux-gnu" 替换为单板的正确target。 # 此target 用于交叉编译。可以通过：code:`gcc -v` 来查询。 target = tvm.target.Target("llvm -device=arm_cpu -mtriple=aarch64-linux-gnu") # 根据设备替换device_key 的值device_key = "rk3399" # 若使用Android 手机，设置use_android 为True use_android = False #### 调优选项#### network = "resnet-18" log_file = "%s.%s.log" % (device_key, network) dtype = "float32" tuning_option = { "log_filename": log_file, "tuner": "xgb", "n_trial": 1500, "early_stopping": 800, "measure_option": autotvm.measure_option( builder=autotvm.LocalBuilder(build_func="ndk" if use_android else "default"), runner=autotvm.RPCRunner( device_key, host="127.0.0.1", port=9190, number=5, timeout=10, ), ), }

最適化を開始する

次に、ネットワークからチューニングタスクを抽出し、チューニングを開始します。その後、シンプルなユーティリティ関数を提供します。これは、タスクリストを順番にチューニングする初期実装です。将来的には、より複雑なチューニングスケジューラを導入する予定です。

 # 可跳过此函数的实现。 def tune_tasks( tasks, measure_option, tuner="xgb", n_trial=1000, early_stopping=None, log_filename="tuning.log", use_transfer_learning=True, ): # 创建tmp 日志文件tmp_log_file = log_filename + ".tmp" if os.path.exists(tmp_log_file): os.remove(tmp_log_file) for i, tsk in enumerate(reversed(tasks)): prefix = "[Task %2d/%2d] " % (i + 1, len(tasks)) # 创建调优器if tuner == "xgb": tuner_obj = XGBTuner(tsk, loss_type="reg") elif tuner == "xgb_knob": tuner_obj = XGBTuner(tsk, loss_type="reg", feature_type="knob") elif tuner == "xgb_itervar": tuner_obj = XGBTuner(tsk, loss_type="reg", feature_type="itervar") elif tuner == "xgb_curve": tuner_obj = XGBTuner(tsk, loss_type="reg", feature_type="curve") elif tuner == "xgb_rank": tuner_obj = XGBTuner(tsk, loss_type="rank") elif tuner == "xgb_rank_knob": tuner_obj = XGBTuner(tsk, loss_type="rank", feature_type="knob") elif tuner == "xgb_rank_itervar": tuner_obj = XGBTuner(tsk, loss_type="rank", feature_type="itervar") elif tuner == "xgb_rank_curve": tuner_obj = XGBTuner(tsk, loss_type="rank", feature_type="curve") elif tuner == "xgb_rank_binary": tuner_obj = XGBTuner(tsk, loss_type="rank-binary") elif tuner == "xgb_rank_binary_knob": tuner_obj = XGBTuner(tsk, loss_type="rank-binary", feature_type="knob") elif tuner == "xgb_rank_binary_itervar": tuner_obj = XGBTuner(tsk, loss_type="rank-binary", feature_type="itervar") elif tuner == "xgb_rank_binary_curve": tuner_obj = XGBTuner(tsk, loss_type="rank-binary", feature_type="curve") elif tuner == "ga": tuner_obj = GATuner(tsk, pop_size=50) elif tuner == "random": tuner_obj = RandomTuner(tsk) elif tuner == "gridsearch": tuner_obj = GridSearchTuner(tsk) else: raise ValueError("Invalid tuner: " + tuner) if use_transfer_learning: if os.path.isfile(tmp_log_file): tuner_obj.load_history(autotvm.record.load_from_file(tmp_log_file)) # 开始调优tsk_trial = min(n_trial, len(tsk.config_space)) tuner_obj.tune( n_trial=tsk_trial, early_stopping=early_stopping, measure_option=measure_option, callbacks=[ autotvm.callback.progress_bar(tsk_trial, prefix=prefix), autotvm.callback.log_to_file(tmp_log_file), ], ) # 选择最佳记录到缓存文件autotvm.record.pick_best(tmp_log_file, log_filename) os.remove(tmp_log_file)

最後に、チューニングタスクが開始され、エンドツーエンドのパフォーマンスが評価されます。

 def tune_and_evaluate(tuning_opt): # 从relay 程序中提取工作负载print("Extract tasks...") mod, params, input_shape, _ = get_network(network, batch_size=1) tasks = autotvm.task.extract_from_program( mod["main"], target=target, params=params, ops=(relay.op.get("nn.conv2d"),) ) # 运行调优任务print("Tuning...") tune_tasks(tasks, **tuning_opt) # 编译具有历史最佳记录的内核with autotvm.apply_history_best(log_file): print("Compile...") with tvm.transform.PassContext(opt_level=3): lib = relay.build_module.build(mod, target=target, params=params) # 导出库tmp = tempdir() if use_android: from tvm.contrib import ndk filename = "net.so" lib.export_library(tmp.relpath(filename), ndk.create_shared) else: filename = "net.tar" lib.export_library(tmp.relpath(filename)) # 上传模块到设备print("Upload...") remote = autotvm.measure.request_remote(device_key, "127.0.0.1", 9190, timeout=10000) remote.upload(tmp.relpath(filename)) rlib = remote.load_module(filename) # 上传参数到设备dev = remote.device(str(target), 0) module = runtime.GraphModule(rlib["default"](dev)) data_tvm = tvm.nd.array((np.random.uniform(size=input_shape)).astype(dtype)) module.set_input("data", data_tvm) # 评估print("Evaluate inference time cost...") print(module.benchmark(dev, number=1, repeat=10)) # 不在网页服务器中运行调优，因为它耗时很久。 # 取消注释运行下一行# tune_and_evaluate(tuning_option)

サンプル出力

最適化には多くのプログラムをコンパイルし、そこから特徴を抽出する必要があるため、高性能なCPUが推奨されます。以下は出力例です。32TBのAMD Ryzen Threadripperデバイスでは、約2時間かかりました。

 Extract tasks... Tuning... [Task 1/12] Current/Best: 22.37/ 52.19 GFLOPS | Progress: (544/1000) | 406.59 s Done. [Task 2/12] Current/Best: 6.51/ 18.77 GFLOPS | Progress: (608/1000) | 325.05 s Done. [Task 3/12] Current/Best: 4.67/ 24.87 GFLOPS | Progress: (480/1000) | 372.31 s Done. [Task 4/12] Current/Best: 11.35/ 46.83 GFLOPS | Progress: (736/1000) | 602.39 s Done. [Task 5/12] Current/Best: 1.01/ 19.80 GFLOPS | Progress: (448/1000) | 262.16 s Done. [Task 6/12] Current/Best: 2.47/ 23.76 GFLOPS | Progress: (672/1000) | 563.85 s Done. [Task 7/12] Current/Best: 14.57/ 33.97 GFLOPS | Progress: (544/1000) | 465.15 s Done. [Task 8/12] Current/Best: 1.13/ 17.65 GFLOPS | Progress: (576/1000) | 365.08 s Done. [Task 9/12] Current/Best: 14.45/ 22.66 GFLOPS | Progress: (928/1000) | 724.25 s Done. [Task 10/12] Current/Best: 3.22/ 15.36 GFLOPS | Progress: (864/1000) | 564.27 s Done. [Task 11/12] Current/Best: 11.03/ 32.23 GFLOPS | Progress: (736/1000) | 635.15 s Done. [Task 12/12] Current/Best: 8.00/ 21.65 GFLOPS | Progress: (1000/1000) | 1111.81 s Done. Compile... Upload... Evaluate inference time cost... Mean inference time (std dev): 162.59 ms (0.06 ms)

Python ソースコードをダウンロード: tune_relay_arm.py

Jupyter Notebook をダウンロード: tune_relay_arm.ipynb

618ZXW