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 下，注