PyTorch Mobile 模型延迟测试

作为一个炼丹科研狗，偶尔总是会有在移动设备上测试模型的需求。目前有不少的移动端推理框架，比如腾讯的NCNN，小米的MACE等等。但是训练模型时一般使用的都是PyTorch框架，利用NCNN这类框架在移动设备上进行推理往往需要将模型转换为对应的格式。这一步一般都是天坑，常常会出现转换之后的模型不能用的情况。如果PyTorch中使用了某些移动端框架没有的算子，那更是天坑中的天坑。

最近PyTorch在移动端也开始发力了，推出了PyTorch Mobile，在IOS和Android设备上都可以很方便将模型转换过去进行推理。因为是PyTorch的亲儿子，所以转换模型十分方便，而且一般也不会出现算子不支持的情况（小声哔哔，我猜的）。本文主要参考Pytorch官方教程 PYTORCH MOBILE PERFORMANCE RECIPES, 在安卓设备上对图片分类、目标检测和语义分割的模型进行了延迟测试，其他模型理论上也可以按照本文的方法进行测试。

环境准备

首先先准备好模型测试所需的所有环境及库文件，以下所有操作均在Ubuntu 20.04上进行。

torch和torchvision
直接按照官网的指引使用pip或者conda安装即可，我安装的版本为torch==1.8.0和torchvision==0.9.0。
Android NDK
因为需要交叉编译，所以需要Android NDK中的arm平台的交叉编译器。在官方的下载页面下载对应平台的NDK文件，我下载的是r21e版本。下载之后解压该文件，并且设置环境变量：

1	export ANDROID_NDK="path/to/android-ndk-r21e"

ADB 调试工具
ADB主要是用来和安卓设设备交互的，比如将本机的文件push到手机。安装过程十分简单，一行命令直接搞定：

1	sudo apt install android-tools-adb

编译模型测试可执行文件

1. Clone PyTorch项目

1	git clone https://github.com/pytorch/pytorch

在大陆地区clone一般速度较慢，可以使用fastgit进行加速，使用命令将所有github的https链接指向fastgit git config --global url."https://hub.fastgit.org/".insteadOf "https://github.com/"

clone完成之后将进入pytorch目录，并且将分支切换到v1.8.0 （这里是为了稳妥起见，我将版本切换为和python中的torch一样的版本）。

1	cd pytorch && git checkout v1.8.0

PyTorch还依赖很多其他第三方项目，运行以下命令将它们全部clone下来：

1	git submodule update --init --recursive

2. 编译speed_benchmark_torch二进制文件

根据官方的教程，一行命令就完事了，编译总共需要几分钟左右：

1	BUILD_PYTORCH_MOBILE=1 ANDROID_ABI=arm64-v8a ./scripts/build_android.sh -DBUILD_BINARY=ON

导出模型

这里以图片分类模型resnet50为例，目标检测和语义分割模型也是大同小异的。这里省略了官方的量化步骤，如果需要量化可以参考官方教程。

import torch
import torch.jit

from torchvision.models import resnet50
from torch.utils.mobile_optimizer import optimize_for_mobile

model = resnet50()
torchscript_model = torch.jit.script(model)
torchscript_model_optimized = optimize_for_mobile(torchscript_model)
torch.jit.save(torchscript_model_optimized, "resnet50.pt")

延迟测试

1. 连接设备

在安卓手机设置中的开发者选项中勾选“USB调试”，然后用数据线将安卓手机连接到电脑，运行命令adb devices，看到类似以下输出则证明已经连接成功了。

1 2	List of devices attached 90e9131c device

2. 推送文件到安卓设备
我们需要将模型测试二进制文件speed_benchmark_torch和模型文件resnet50推送到/data/local/tmp目录下面。

1 2	adb push <speedbenchmark-torch> /data/local/tmp adb push <path-to-scripted-model> /data/local/tmp

3. 进行测试

使用命令adb shell进行安卓命令行，然后进入目录cd /data/local/tmp，赋予可执行权限chmod a+x speed_benchmark_torch。最后运行以下命令测试resnet50的延迟：

1 2	speed_benchmark_torch --model=resnet50.pt \ --input_dims="1,3,224,224" --input_type="float"

以下为终端输出：

Starting benchmark.
Running warmup runs.
Main runs.
Main run finished. Microseconds per iter: 172630. Iters per second: 5.79275

其时间单位是微秒，所以每个batch推理所需的时间为172.63毫秒。另外，还可以通过调整input_dims参数控制batch size和输入分辨率。

此外，我还测试了分割模型torchvision.models.segmentation.lraspp_mobilenet_v3_large的延迟：

1 2	speed_benchmark_torch --model=lraspp_mobilenet_v3_large.pt \ --input_dims="1,3,512,1024" --input_type="float"

终端输出为：

Starting benchmark.
Running warmup runs.
Main runs.
Main run finished. Microseconds per iter: 547199. Iters per second: 1.82749