Pytorch模型转Android端模型

Posted AI浩

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了Pytorch模型转Android端模型相关的知识,希望对你有一定的参考价值。

android端部署pytorch模型,需要将模型转为.pt文件,转换方法如下:

import torch
import torch.utils.data.distributed

# pytorch环境中
model_pth = 'model_31_0.96.pth' # resnet18模型的参数文件
mobile_pt ='model.pt' # 将resnet18模型保存为Android可以调用的文件

model = torch.load(model_pth)
model.eval() # 模型设为评估模式
device = torch.device('cpu')
model.to(device)
# 1张3通道224*224的图片
input_tensor = torch.rand(1, 3, 224, 224) # 设定输入数据格式

mobile = torch.jit.trace(model, input_tensor) # 模型转化
mobile.save(mobile_pt) # 保存文件

以上是关于Pytorch模型转Android端模型的主要内容,如果未能解决你的问题,请参考以下文章

[Pytorch]Pytorch 保存模型与加载模型(转)

pytorch模型转caffe模型

模型转换 PyTorch转ONNX 入门

SwinTransformer模型转化:pytorch模型转keras。

pytorch模型压缩方法总结

pytorch模型文件pth详解