# PyTorch Mobile 模型延迟测试
作为一个炼丹科研狗,偶尔总是会有在移动设备上测试模型的需求。目前有不少的移动端推理框架,比如腾讯的NCNN,小米的MACE等等。但是训练模型时一般使用的都是PyTorch框架,利用NCNN这类框架在移动设备上进行推理往往需要将模型转换为对应的格式。这一步一般都是天坑,常常会出现转换之后的模型不能用的情况。如果PyTorch中使用了某些移动端框架没有的算子,那更是天坑中的天坑。
最近PyTorch在移动端也开始发力了,推出了PyTorch Mobile,在IOS和Android设备上都可以很方便将模型转换过去进行推理。因为是PyTorch的亲儿子,所以转换模型十分方便,而且一般也不会出现算子不支持的情况(小声哔哔,我猜的)。本文主要参考Pytorch官方教程 PYTORCH MOBILE PERFORMANCE RECIPES (opens new window), 在安卓设备上对图片分类、目标检测和语义分割的模型进行了延迟测试,其他模型理论上也可以按照本文的方法进行测试。
# 环境准备
首先先准备好模型测试所需的所有环境及库文件,以下所有操作均在Ubuntu 20.04上进行。
torch和torchvision直接按照官网的指引使用pip或者conda安装即可,我安装的版本为torch==1.8.0和torchvision==0.9.0。Android NDK因为需要交叉编译,所以需要Android NDK中的arm平台的交叉编译器。在官方的下载页面 (opens new window)下载对应平台的NDK文件,我下载的是r21e版本。下载之后解压该文件,并且设置环境变量:
export ANDROID_NDK="path/to/android-ndk-r21e"
ADB调试工具 ADB主要是用来和安卓设设备交互的,比如将本机的文件push到手机。安装过程十分简单,一行命令直接搞定:
sudo apt install android-tools-adb
# 编译模型测试可执行文件
1. Clone PyTorch项目
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一样的版本)。
cd pytorch && git checkout v1.8.0
PyTorch还依赖很多其他第三方项目,运行以下命令将它们全部clone下来:
git submodule update --init --recursive
2. 编译speed_benchmark_torch二进制文件
根据官方的教程,一行命令就完事了,编译总共需要几分钟左右:
BUILD_PYTORCH_MOBILE=1 ANDROID_ABI=arm64-v8a ./scripts/build_android.sh -DBUILD_BINARY=ON
# 导出模型
这里以图片分类模型resnet50为例,目标检测和语义分割模型也是大同小异的。这里省略了官方的量化步骤,如果需要量化可以参考官方教程 (opens new window)。
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")
2
3
4
5
6
7
8
9
10
# 延迟测试
1. 连接设备
在安卓手机设置中的开发者选项中勾选“USB调试”,然后用数据线将安卓手机连接到电脑,运行命令adb devices,看到类似以下输出则证明已经连接成功了。
List of devices attached
90e9131c device
2
2. 推送文件到安卓设备
我们需要将模型测试二进制文件speed_benchmark_torch和模型文件resnet50推送到/data/local/tmp目录下面。
adb push <speedbenchmark-torch> /data/local/tmp
adb push <path-to-scripted-model> /data/local/tmp
2
3. 进行测试
使用命令adb shell进行安卓命令行,然后进入目录cd /data/local/tmp,赋予可执行权限chmod a+x speed_benchmark_torch。最后运行以下命令测试resnet50的延迟:
speed_benchmark_torch --model=resnet50.pt \
--input_dims="1,3,224,224" --input_type="float"
2
以下为终端输出:
Starting benchmark.
Running warmup runs.
Main runs.
Main run finished. Microseconds per iter: 172630. Iters per second: 5.79275
2
3
4
其时间单位是微秒,所以每个batch推理所需的时间为172.63毫秒。另外,还可以通过调整input_dims参数控制batch size和输入分辨率。
此外,我还测试了分割模型torchvision.models.segmentation.lraspp_mobilenet_v3_large的延迟:
speed_benchmark_torch --model=lraspp_mobilenet_v3_large.pt \
--input_dims="1,3,512,1024" --input_type="float"
2
终端输出为:
Starting benchmark.
Running warmup runs.
Main runs.
Main run finished. Microseconds per iter: 547199. Iters per second: 1.82749
2
3
4