ARM CPU神经网络自动调度

最新推荐文章于 2023-05-11 18:12:47 发布

wujianming_110117

最新推荐文章于 2023-05-11 18:12:47 发布

阅读量542

点赞数

CC 4.0 BY-SA版权

分类专栏：集成电路大数据智能芯片

吴建明

本文链接：https://blog.csdn.net/wujianing_110117/article/details/115192279

ARM CPU神经网络自动调度
对特定设备和工作负载进行自动调整对于获得最佳性能至关重要。这是一个有关如何通过RPC使用自动调度器为ARM CPU调整整个神经网络的教程。
为了自动调整神经网络，将网络划分为小的子图，对其进行独立调整。每个子图被视为一个搜索任务。任务调度程序可以对时间进行分片，为这些任务动态分配时间资源。任务调度程序可以预测每个任务对端到端执行时间的影响，优先安排可以最大程度地减少执行时间的任务。
对于每个子图，使用compute声明tvm/python/topi，获取张量表达式形式的计算DAG。然后，使用自动调度器来构造此DAG的搜索空间，搜索良好的调度（低级优化）。
与依靠手动模板定义搜索空间的基于模板的autotvm不同，自动调度程序不需要任何调度模板。换句话说，自动调度程序仅在其中使用计算声明，tvm/python/topi而不使用现有的调度模板。
本文无法在Windows或最新版本的macOS上运行。要使其运行，需要将本文的内容包装在一个if name == “main”:块中。
import numpy as np

import tvm
from tvm import relay, auto_scheduler
import tvm.relay.testing
from tvm.contrib import graph_runtime
from tvm.contrib.utils import tempdir
定义网络
首先，需要使用中继前端API定义网络。可以加载一些tvm.relay.testing预定义的网络。可以从MXNet，ONNX，PyTorch和TensorFlow加载模型。
对于卷积神经网络，尽管自动调度程序可以在任何布局下正常工作，发现使用NHWC布局通常可以实现最佳性能。使用自动调度程序对NHWC布局实施了更多优化。建议将模型转换为NHWC布局以使用自动调度程序。可以使用ConvertLayout传递在TVM中进行布局转换。
def get_network(name, batch_size, layout=“NHWC”, dtype=“float32”):
“”“Get the symbol definition and random weight of a network”""

# auto-scheduler prefers NHWC layout
if layout == "NHWC":
    image_shape = (224, 224, 3)
elif layout == "NCHW":
    image_shape = (3, 224, 224)
else:
    raise ValueError("Invalid layout: " + layout)

input_shape = (batch_size,) + image_shape
output_shape = (batch_size, 1000)

if name.startswith("resnet-"):
    n_layer = int(name.split("-")[1])
    mod, params = relay.testing.resnet.get_workload(
        num_layers=n_layer,
        batch_size=batch_size,
        layout=layout,
        dtype=dtype,
        image_shape=image_shape,
    )
elif name.startswith("resnet3d-"):
    n_layer = int(name.split("-")[1])
    mod, params = relay.testing.resnet.get_workload(
        num_layers=n_layer,
        batch_size=batch_size,
        layout=layout,
        dtype=dtype,
        image_shape=image_shape,
    )
elif name == "mobilenet":
    mod, params = relay.testing.mobilenet.get_workload(
        batch_size=batch_size, layout=layout, dtype=dtype, image_shape=image_shape
    )
elif name == "squeezenet_v1.1":
    assert layout == "NCHW", "squeezenet_v1.1 only supports NCHW layout"
    mod, params = relay.testing.squeezenet.get_workload(
        version="1.1",
        batch_size=batch_size,
        dtype=dtype,
        image_shape=image_shape,
    )
elif name == "inception_v3":
    input_shape = (batch_size, 3, 299, 299) if layout == "NCHW" else (batch_size, 299, 299, 3)
    mod, params = relay.testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype)
elif name == "mxnet":
    # an example for mxnet model
    from mxnet.gluon.model_zoo.vision import get_model

    assert layout == "NCHW"

    block = get_model("resnet50_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)

return mod, params, input_shape, output_shape

启动RPC跟踪器
TVM使用RPC会话与ARM板进行通信。在调整期间，调谐器会将生成的代码发送到电路板上，测量电路板上的代码速度。
为了扩大调整范围，TVM使用RPC Tracker来管理分布式设备。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跟踪器
将设备注册到跟踪器。第一步是为ARM设备构建TVM运行时。
• 对于Linux：遵循本节在设备上构建TVM运行时，在设备上构建TVM运行时。然后通过以下方式将设备注册到跟踪器
• python -m tvm.exec.rpc_server --tracker=[HOST_IP]:9190 --key=rasp4b-64
（替换[HOST_IP]为主机的IP地址）
• 对于Android：按照以下Readme在Android设备上安装TVM RPC APK。确保可以通过android rpc测试。然后，已经注册了设备。在调整过程中，必须转到开发人员选项，启用“更改时保持屏幕唤醒”并为手机充电以使其稳定。
注册设备后，通过查询rpc_tracker进行确认
python -m tvm.exec.query_rpc_tracker --host=0.0.0.0 --port=9190
例如，如果有2个Huawei mate10 pro，11个具有64位操作系统的Raspberry Pi 4B和2个rk3399，输出可以是
Queue Status

key total free pending

mate10pro 2 2 0
rk3399 2 2 0
rasp4b-64 11 11 0

可以将多个设备注册到跟踪器，加快调谐过程中的测量速度。
设置调整选项
调整之前，应该应用一些配置。在这里，带有64位操作系统（Ubuntu 20.04）的Raspberry Pi 4b 4GB主板为例。在设置中，应该相应地修改use_ndk目标和device_key。如果使用的是Android手机，设置为True。