目录
- 环境
- 获取demo
- 将pt文件导出为ptl文件
- 修改demo
- 修改PrePostProcessor
- 增加ptl文件并增加类别文件
- 修改MainActivity
- 大功告成
环境
Ubuntu22.10
Pytorch2.0.1+cu117
Android Studio Flamingo | 2022.2.1 Patch 1
获取demo
git clone https://github.com/pytorch/android-demo-app.git
打开android-demo-app/ObjectDectection项目。
打开app/build.gradle文件,修改pytorch版本
这里遇到最大的坑,就是训练的pytorch版本和android的包版本相差过大,网上大多数安卓这里都是1.10.0
implementation 'org.pytorch:pytorch_android_lite:1.13.1'
implementation 'org.pytorch:pytorch_android_torchvision_lite:1.13.1'
将pt文件导出为ptl文件
修改yolov5中的export.py,其中export_torchscript函数不能导出ptl文件,修改为以下:
def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
# YOLOv5 TorchScript model export
try:
print(f'\n{prefix} starting export with torch {torch.__version__}...')
f = file.with_suffix('.torchscript.pt')
f = str(f)
fl = file.with_suffix('.torchscript.ptl')
ts = torch.jit.trace(model, im, strict=False)
(optimize_for_mobile(ts) if optimize else ts).save(f)
(optimize_for_mobile(ts) if optimize else ts)._save_for_lite_interpreter(str(fl))
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
print(f'{prefix} export success, saved as {fl} ({file_size(fl):.1f} MB)')
except Exception as e:
print(f'{prefix} export failure: {e}')
return fl, None
运行导出ptl文件
python export.py --weights <your pt file path> --include torchscript
在pt文件的目录下会发现这样的文件,第二个最后后缀为ptl的是我们所需要的文件
修改demo
修改PrePostProcessor
在PrePostProcessor文件中将mOutputColumn修改为网络模型类别数+5,我这里网络模型类别数是2,所以写7
增加ptl文件并增加类别文件
在assets文件夹中放入之前导出的ptl文件并且添加一个类别文件,类别文件中填写网络模型类别的name,以回车间隔,如下图
修改MainActivity
这里修改为ptl文件名和classes文件名