AI挑战营(进阶) 三：获取onnx模型

zhuxirui

AI挑战营(进阶) 三：获取onnx模型 [复制链接]

紧跟上节我们介绍了InsightFace中各个流程使用到的模型并选择了retinaface和arcface来实现我们人脸识别的任务。

 

根据LuckFox官方和RKNPU2的文档指引，我们需要先获取onnx（最常用）的模型格式，但这在InsightFace官方程序中并未提供，我们可以自己找一下已有的insighface，onnx化部署工作。

一般我会先在onnx_model_zoo项目下寻找，这个项目由ONNX官方收录了许多完成了预训练的ONNX模型，可以降低我们的转换工作的难度。

 

但是很不幸，这里只提供了retinanet的网络模型，并没有retinaface的模型，而且里面的arcface模型采用的骨干网络为ResNet100过于复杂庞大，虽然也可以转换为rknn而且精度更高，但不适合在RV1106这样计算资源十分紧张的端侧设备上部署，我们希望能找到一个使用mobilenetssd作为骨干网络的arcface模型。

 

最后我在这两个项目里面找到了我们需要的onnx模型。

RetinaFace:

链接已隐藏，如需查看请登录或者注册

ArcFace:

链接已隐藏，如需查看请登录或者注册

其中的retinaface模型仍需要我们从pth转换到onnx但这个过程比较简单，前面也已经有网友给出了转换步骤和脚本

 

git clone https://github.com/bubbliiiing/retinaface-pytorch
cd retinaface-pytorch
touch 2onnx.py

 

#2onnx.py
from nets.retinaface import RetinaFace
from utils.config import cfg_mnet
import torch
model_path='model_data/Retinaface_mobilenet0.25.pth' #模型路径
model=RetinaFace(cfg=cfg_mnet,pretrained = False) #模型初始化
device = torch.device('cpu')
model.load_state_dict(torch.load(model_path,map_location=device),strict=False) #模型加载
net=model.eval()
example=torch.rand(1,3,640,640) #给定输入
torch.onnx.export(model,(example),'model_data/retinaface.onnx',verbose=True,opset_version=9) #导出

 

接着我们可以通过https://netron.app/，将模型拖入观察一下模型结构，主要是确认一下输入和输出的大小

Retianface:

Arcface:

我们发现arcface模型的batch没有固定，rv1106似乎也不支持批量推理和动态输入，所以需要我们手动调整一下

touch crop_model.py

 

#Crop_model:
import onnx
def change_input_size(model_path, new_size1):
model = onnx.load(model_path)
for input in model.graph.input:
input.type.tensor_type.shape.dim[0].dim_value = new_size1
onnx.save(model, model_path)
change_input_size("w600k_mbf.onnx", 1)

 

调整完成后得到的arcface模型结构如下：

至此我们就得到了retinaface和arcface两个关键的onnx模型，下一节我们将完成onnx2rknn的工作。

这里也给出了onnx模型供大家参考