在x86的Docker中构建TVM的ARM环境

Posted 夏小悠

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了在x86的Docker中构建TVM的ARM环境相关的知识,希望对你有一定的参考价值。

文章目录

前言

  本篇文章介绍如何在x86docker中构建tvmARM环境,以及如何使用RPC使编译在x86的环境中运行在arm的环境中,还介绍了如何在arm环境中进行编译和运行,并提供详细的示例以供验证,其中包括rpc测试代码,acl测试代码,pytorch模型在arm上的推理以及在arm上进行autotvm
  如下图所示,显示的是x86架构的cpu信息:

  强烈建议使用ubuntu:20.04这个版本,ubuntu:18.04这个版本在升级glibc时会掉坑里!!!

1. 加载arm-ubuntu镜像

  从docker镜像库中拉取arm-ubuntu镜像:

docker pull arm64v8/ubuntu:20.04

  由于本地的cpux86架构,没办法直接运行arm架构的镜像,需要借助第三方工具:QEMU
  QEMU是一个通用的开源的跨平台仿真模拟器,可以模拟在特定的体系结构下的应用的执行或者构建,比如在x86的体系结构的操作系统上运行ARM的应用。
  目前看到使用qemu进行模拟的有两种方式:一种是结合docker使用[本博客使用的],另一种是使用qemu官方的源码进行编译,手动安装相应系统的iso,可以参考这篇博客

  使用docker构建的arm环境,用lscpu指令查看cpu时,cpumodel name仍然是intel的,archaarch64;使用源码编译成的qemu-system-aarch64构建模拟器时会指定具体的cpu型号,比如qemu-system-aarch64 -cpu cortex-a72,由于没有尝试这种方式,因此不确定cpumodel nameintel的,还是arm的。

docker run --rm --privileged multiarch/qemu-user-static --reset -p yes

  这条指令会安装qemu-user-static,安装完毕后就可以正常在x86docker中运行ARM架构下的ubuntu镜像了,不过这条指令还没有测试。我使用的是下面的方法,可以参考这篇文章
  先下载qemu-aarch64-static安装包:


  对qemu-aarch64-static进行配置:

sudo cp qemu-aarch64-static /usr/bin/
sudo chmod +x /usr/bin/qemu-aarch64-static
# 注册QEMU虚拟机
docker run --rm --privileged multiarch/qemu-user-static:register

  然后就可以正常的加载arm-ubuntu了:

docker run --platform linux/arm64/v8 -it -v /home/liyanpeng/arm64v8_work:/home/liyanpeng/arm64v8_work -w /home/liyanpeng arm64v8/ubuntu:20.04 bash

# uname -a
# lscpu

2. 安装acl库

  ARM 计算库(Arm Compute Library, ACL),是为 ARM 架构的 CPUGPU 提供加速内核的开源项目。可以从 ARM-software 下载预构建的二进制文件:

# 将压缩包解压到 acl_tmp 目录
tar -zxvf arm_compute-v22.08-bin-linux-arm64-v8.2-a-neon.tar.gz -C acl_tmp

  直接编译ARM架构下的runtime时会报错,需要手动调整相应的目录,可以参考tvm官方给的一个脚本:ubuntu_download_arm_compute_lib_binaries.sh

cp -r acl_tmp/include acl/
cp -r acl_tmp/arm_compute acl/include/
cp -r acl_tmp/support acl/include/
cp -r acl_tmp/utils acl/include/
cp -r acl_tmp/lib/arm64-v8.2-a-neon acl/lib

3. 编译arm运行时

  在编译之前,仍然需要在arm-ubuntu配置基本的环境,包括C/C++CMakePython等基本环境,具体可以参考我的这篇文章:《tvm在linux环境下的安装与编译及vscode如何配置tvm的远程连接调试环境》,这里不在赘述。

  修改build/config.cmake文件:

set(USE_LLVM OFF)	# line 136(default)
set(USE_ARM_COMPUTE_LIB OFF)	# line 236(default)
set(USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR "/home/liyanpeng/arm64v8_work/acl")	# acl的路径

  进行编译:

cd build
cmake ..
make runtime -j6

  编译成功后的信息如下:



  不要忘了添加tvmpython环境:

export PYTHONPATH=$PYTHONPATH:/home/liyanpeng/arm64v8_work/tvm_work/tvm/python

  tvm版本验证:

import tvm

print(tvm.__version__)

4. 编译在x86运行在arm

4.1 在x86的环境中构建arm的编译环境

  修改build/config.cmake文件:

set(USE_LLVM ON)	# line 136
set(USE_ARM_COMPUTE_LIB ON)	# line 236
set(USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR OFF) # line 237

  进行编译:

cd build
cmake ..
make -j6

  很快就编译好了:


  构建完成后即可在x86的环境中编译arm支持的算子,仅仅是编译,在x86上是不能运行直接运行的。

4.2 测试x86-ubuntu与arm-ubuntu能否ping通

  可以借助RPC(Remote Produce Call)来实现 编译在x86,运行在ARM,因此,需要知道arm-ubuntuip地址:
  在arm-ubuntu中安装网络工具包:

apt-get update
# ifconfig
apt-get install net-tools
# ping
apt-get install inetutils-ping

  查看arm-ubuntuip地址:

  查看x86-ubuntuip地址:


  测试x86-ubuntuarm-ubuntu能否ping通:

# x86-ubuntu
ping 172.17.0.2

# arm-ubuntu
ping 172.17.0.3

4.3 调用RPC

  在arm-ubuntu环境中启动RPC

python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090

  启动成功的信息如下:

  在x86-ubuntu环境中创建rpc_test.py文件,内容如下:

# rpc_test.py
import numpy as np

import tvm
from tvm import te
from tvm import rpc
from tvm.contrib import utils, tar


n = tvm.runtime.convert(1024)
A = te.placeholder((n,), name="A")
B = te.compute((n,), lambda i: A[i] + 1.0, name="B")
s = te.create_schedule(B.op)

local_demo = False

if local_demo:
    target = "llvm"
else:
    # target = "llvm -mtriple=armv7l-linux-gnueabihf"     # Raspberry Pi 3B
    # target = "llvm -mtriple=aarch64-linux-gnu"
    # target = tvm.target.arm_cpu() # error: error adding symbols: file in wrong format
    target = "llvm -mtriple=aarch64-linux-gnu -mattr=+neon"

func = tvm.build(s, [A, B], target=target, name="add_one")

# save the lib at a local temp folder
temp = utils.tempdir()
path = temp.relpath("lib_rpc_test.tar")
func.export_library(path, tar.tar)

print("lib path: ", path)

if local_demo:
    remote = rpc.LocalSession()
else:
    # The following is my environment, change this to the IP address of your target device
    host = "172.17.0.5"		# arm-ubuntu ip
    port = 9090
    remote = rpc.connect(host, port)

remote.upload(path)
func = remote.load_module("lib_rpc_test.tar")

# create arrays on the remote device
dev = remote.cpu()
a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), dev)
b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), dev)

# the function will run on the remote device
func(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)

time_f = func.time_evaluator(func.entry_name, dev, number=10)
cost = time_f(a, b).mean
print("%g secs/op" % cost)

  上面代码演示了一个加法运算,执行结果如下:

  在arm-ubuntu可以看到来自x86-ubuntu的连接信息:

4.4 ACL的使用

  ACL的使用可以参看tvm官方的示例文档,这里给出了例子,使用方式同上小节的一样用的是RPC

# acl_test.py
import tvm
from tvm import relay
from tvm import rpc
from tvm.contrib import utils, tar
from tvm.relay.op.contrib.arm_compute_lib import partition_for_arm_compute_lib

import numpy as np


data_type = "float32"
data_shape = (1, 14, 14, 512)
strides = (2, 2)
padding = (0, 0, 0, 0)
pool_size = (2, 2)
layout = "NHWC"
output_shape = (1, 7, 7, 512)

# use a single max_pool2d operator
data = relay.var('data', shape=data_shape, dtype=data_type)
out = relay.nn.max_pool2d(data, pool_size=pool_size, strides=strides, layout=layout, padding=padding)
module = tvm.IRModule.from_expr(out)

# annotate and partition the graph for ACL
module = partition_for_arm_compute_lib(module)

# build the Relay graph.
target = "llvm -mtriple=aarch64-linux-gnu -mattr=+neon"
with tvm.transform.PassContext(opt_level=3, disabled_pass=["AlterOpLayout"]):
    lib = relay.build(module, target=target)

# export the module
lib_path = './lib_acl.tar'
# cross_compile = 'aarch64-linux-gnu-c++'
# lib.export_library(lib_path, cc=cross_compile)
lib.export_library(lib_path)

# rpc
host = "172.17.0.2"		# arm-ubuntu ip
port = 9090
remote = rpc.connect(host, port)

remote.upload(lib_path)
loaded_lib = remote.load_module("lib_acl.tar")

# run Inference
# dev = tvm.cpu(0)
# loaded_lib = tvm.runtime.load_module('lib_acl.so')
dev = remote.cpu(0)
module = tvm.contrib.graph_executor.GraphModule(loaded_lib['default'](dev))
d_data = np.random.uniform(0, 1, data_shape).astype(data_type)
map_inputs = 'data': d_data
module.set_input(**map_inputs)
module.run()

# get output
output = module.get_output(0)
print("TVM MaxPool2d[acl] output: ", output)

  运行结果如下:

  上面的示例仅显示了ACL如何用于单个Maxpool2D的基本示例。如果想看到网络中每个算子的实现,请参阅tests/python/contrib/test_arm_compute_lib

5. arm版的tvm编译和运行时环境

5.1 构建arm版的tvm编译和运行时环境

  ARM版的tvm编译和运行时环境的构建同x86版的几乎一样,可以按照这篇文章:《tvm在linux环境下的安装与编译及vscode如何配置tvm的远程连接调试环境》进行配置,这里不在赘述。不过需要注意的是,在arm-ubuntu中并没有找到conda/build-environment.yaml文件中要求的llvmdev ==10.0.0这个版本,因此这里略作修改:

# conda/build-environment.yaml
# 这里将llvmdev更改为10.0.1版本
# 这样在编译时cmake会自动安装llvm
llvmdev ==10.0.1

  基于上述配置,再次修改build/config.cmake文件:

set(USE_LLVM ON)	# line 136
set(USE_ARM_COMPUTE_LIB ON)	# line 236
set(USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR "/home/liyanpeng/arm64v8_work/acl") # line 237

  然后进行编译:

cd build
cmake ..
make -j6

  俩小时过去了。。。

  编译成功后的信息如下:

5.2 关于ubuntu 18.04升级glibc掉入坑里这件事

  遗憾的是,在pytorch 1.7.1版本及以前的版本中,官方并没有提供arm版本的pytorch,这里有两种解决方式:
  (1) 从社区下载非官方版本的pytorch-aarch64,比如:KumaTea
  (2) 官方从pytorch 1.8.0版本开始提供了arm版本的pytorch,可以选择更高版本的pytorch,不过还是要说一下,tvm官方目前[文章发布时]支持pytorch 1.7pytorch 1.4两个大版本,其他版本可能不稳定。
  这里选择了非官方版本的pytorch-aarch64

pip install torch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 -f https://torch.kmtea.eu/whl/stable-cn.html

  在查看pytorch版本时报错:ImportError: /lib/aarch64-linux-gnu/libc.so.6: version "GLIBC_2.28" not found[最初用的是ubuntu:18.04这个版本,换成ubuntu:20.04这个版本不会报错,可直接看5.3小节]


  查看系统当前glibc版本

ldd --version
# or
strings /lib/aarch64-linux-gnu/libm.so.6 | grep GLIBC_

# Ubuntu 18.04: 2.27
# Ubuntu 20.04: 2.31


  好了,可以打住了,建议掉头,直接跳到5.3小节,不然等会儿掉坑里可能出不来!!!

  解决方法可以参考这篇博客

# 安装依赖
apt-get install gawk
apt-get install bison

apt-get install wget

# 下载、解压并配置
wget http://ftp.gnu.org/gnu/libc/glibc-2.28.tar.gz
tar -zxvf glibc-2.28.tar.gz
cd glibc-2.28
mkdir build
cd build
../configure --prefix=/usr/local --disable-sanity-checks

# 安装
make -j6
make install

  安装过程中的一些日志信息如下:


  没有任何报错信息,说明安装成功。

# 查看原始的软连接
ll /lib/aarch64-linux-gnu/libc.so.6

  按网上一些教程出现Segmentation fault错误,导致常用的lscpclear这些指令无法使用,解决方法:

# export LD_PRELOAD=/lib/aarch64-linux-gnu/libc-2.27.so:/lib/aarch64-linux-gnu/ld-2.27.so
unset LD_PRELOAD
# 取消软连接
LD_PRELOAD=/lib/aarch64-linux-gnu/libc-2.27.so unlink /lib/aarch64-linux-gnu/libc.so.6
# 重新恢复
LD_PRELOAD=/lib/aarch64-linux-gnu/libc-2.27.so ln -s /lib/aarch64-linux-gnu/libc-2.27.so /lib/aarch64-linux-gnu/libc.so.6

  建立软连接:

# 复制 libc
cp /usr/local/lib/libc-2.28.so /lib/aarch64-linux-gnu/
cp /usr/local/lib/ld-2.28.so /lib/aarch64-linux-gnu/

cd /lib/aarch64-linux-gnu/
# ll ld-linux-aarch64.so.1
# ll libc.so.6

ln -sf /lib/aarch64-linux-gnu/libc-2.28.so /lib/aarch64-linux-gnu/libm.so.6
# 无效, 仍然是2.27版本

  对比了aarch64-linux-gnu目录与glibc-2.28的安装目录,发现好多库名一样,只是版本号不一样,是不是都要进行替换???

5.3 验证安装是否成功

  到了这一步,如果ubuntu 18.04glibc没有升级成功,那就使用ubuntu 20.04吧,以下是在arm-ubuntu 20.04中进行的,查看pytorch版本:

  pytorch模型验证:

# from_pytorch.py

import tvm
from tvm import relay
from tvm.contrib.download import download_testdata
import numpy as np

import torch
import torchvision

######################################################################
# Load a pretrained PyTorch model
pth_file = 'resnet18-f37072fd.pth'
model = torchvision.models.resnet18()
ckpt = torch.load(pth_file)
model.load_state_dict(ckpt)
model = model.eval()

# We grab the TorchScripted model via tracing
input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model, input_data).eval()

######################################################################
# Load a test image
from PIL import Image

img_path = 'cat.png'
img = Image.open(img_path).resize((224, 224))

# Preprocess the image and convert to tensor
from torchvision import transforms

my_preprocess = transforms.Compose(
    [
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)
img = my_preprocess(img)
img = np.expand_dims(img, 0)

######################################################################
# Import the graph to Relay
input_name = "input0"
shape_list = [(input_name, img.shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)

######################################################################
# Relay Build
target = tvm.target.arm_cpu()
dev = tvm.cpu(0)
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target=target, params=params)

######################################################################
# Execute the portable graph on TVM
from tvm.contrib import graph_executor

dtype = "float32"
m = graph_executor.GraphModule(lib["default"](dev))
# Set inputs
m.set_input(input_name, tvm.nd.array(img.astype(dtype)))
# Execute
m.run()
# Get outputs
tvm_output = m.get_output(0)

#####################################################################
# Look up synset name
synset_path = 'imagenet_synsets.txt'
with open(synset_path) as f:
    synsets = f.readlines()

synsets = [x.strip() for x in synsets]
splits = [line.split(" ") for line in synsets]
key_to_classname = spl[0]: " ".join(spl[1:]) for spl in splits

class_path = 'imagenet_classes.txt'
with open(class_path) as f:
    class_id_to_key = f.readlines()

class_id_to_key = [x.strip() for x in class_id_to_key]

# Get top-1 result for TVM
top1_tvm = np.argmax(tvm_output.numpy()[0])
tvm_class_key = class_id_to_key[top1_tvm]

# Convert input to PyTorch variable and get PyTorch result for comparison
with torch.no_grad():
    torch_img = torch.from_numpy(img)
    output = model(torch_img)

    # Get top-1 result for PyTorch
    top1_torch = np.argmax(output.numpy())
    torch_class_key = class_id_to_key[top1_torch]

print("Relay top-1 id: , class name: ".format(top1_tvm, key_to_classname[tvm_class_key]))
print("Torch top-1 id: , class name: ".format(top1_torch, key_to_classname[torch_class_key]))

  验证结果如下【速度不是有一点慢】:


  autotvm也是ok的:

  但感觉有些问题,在x86task只有13[resnet18],而在arm上确是26个,经查验这26task即包含arm的,还包含x86,就很奇怪,这个问题还待解决:

结束语

  本篇文章是对在x86docker中搭建tvmarm环境的一次尝试,由于服务器上使用的是ubuntu 18.04,本想着在本地模拟一个同服务器相同的版本,结果遇到了ubuntu 18.04升级glibc这个坑,暂时还没有填上。不过,成功在x86docker中搭建了tvmarm环境,并能进行编译和运行。
  tvm仍在持续学习中,欢迎评论区交流哦!!!

以上是关于在x86的Docker中构建TVM的ARM环境的主要内容,如果未能解决你的问题,请参考以下文章

使用 docker buildx 构建多 CPU 架构镜像

x86机器上运行arm64 docker

x86机器上运行arm64 docker

arm搭建x86运行时

qemu-user-static, 在x86机器上编译arm架构docker

docker镜像能跨平台运行吗?(可以跨平台,但是镜像架构必须与操作系统底层架构相同)(x86架构arm架构)