pytorch模型部署到Android

pytorch版本：1.6.0
pytorch-android版本：1.6.0

1 model.pt->model-script.pt

若模型上一次由GPU训练得到，需要转换成CPU形式

import torch

device = torch.device('cpu')
net=torch.load('model.pt', map_location = device)
torch.save(net,'model-cpu.pt')

然后把model.pt转换成Pytorch-script，以便在安卓上运行

import torch

# 如果网络使用的是class Net的定义法，而不是快速搭建法，需要在此处引入class Net的定义

model = torch.load("model-cpu.pt")
 
model.eval()
input_tensor = torch.rand(1,100) # 这里写你的模型的输入张量形状
# 笔者的模型为输入一维张量，100 features，后文的输入尺寸也与此对应
 
script_model = torch.jit.trace(model,input_tensor)
script_model.save("model-script.pt")

2 Android Studio 配置

新建一个 C++ Native 项目，选择c++11

在 build.gradle (Module) repositories 下，注释掉 jcenter()，添加镜像源：

maven{ url 'http://maven.aliyun.com/nexus/content/repositories/central/'}
maven{ url 'http://maven.aliyun.com/nexus/content/repositories/jcenter'}

在 build.gradle (App) dependencies 下，添加依赖：

implementation 'org.pytorch:pytorch_android:1.6.0'
implementation 'org.pytorch:pytorch_android_torchvision:1.6.0'

添加后，Android Studio 提示同步，点击 Sync Now，开始同步Gradle

在main目录下新建assets目录，把之前转换好的model-script.pt放到该目录下

3 调用模型

在 Activity 中，添加：

float[] data=new float[100];

//do something，为data赋值，可以是从文件加载、从用户输入、从相机获取，等等

copyAssetAndWrite("model-script.pt");//把模型从assets写入缓存，以便调用

Module module = Module.load(getCacheDir()+"/model-script.pt");//从缓存区加载模型

long shape[]={1,100};//模型输入形状

Tensor tensor=Tensor.fromBlob(data,shape);//tensor初始化方法

IValue input=IValue.from(tensor);
Tensor output=module.forward(input).toTensor();
float predict[]=output.getDataAsFloatArray();

这样，predict[]就是从网络得到的输出了

4 Build Apk,完成