pytorch导入Android环境中

最近有需求就是需要把用pytorch搭建的网络导入到Android手机上，意思是做一个app来展示网络的功能。我发现，关于这种的博客很少，所以就写一个博客来记录我的过程。

举个例子

下面的就是最基本的，直接使用训练好的模型resnet18，这个就是在ImageNet上训练的模型

import torch
import torchvision

model = torchvision.models.resnet18(pretrained=True)
model.eval()  
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("D:\CameraDemo/app/src/main/assets/modell.pt")

其他具体的在AndroidStudio上的操作
比如什么加载模型
可以参考这篇博客：

我用的是python3.8，pytorch是1.6.0
所以在导入库的时候，要导入1.6 这个具体要看你的版本是啥

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

那么对于我们自定义的数据集怎么操作呢

下面的ResNet18就是我们自己写的网络，用来训练自定义数据集，然后导入参数之后就可以保存到AndroidStudio的工程目录下面了。

import torch
from ResNet import ResNet18

model = ResNet18(5)