AMB82 MobileFaceNet Convert To ONNX

MobileFaceNet

MobileFaceNet 是一种高效的Convolutional Neural Network (CNN) 模型,它使用超过 100 万个参数。
MobileFaceNet 用于特征提取。 由于 MobileFaceNet 是轻量级模型类型之一,我们可以将此人脸识别系统应用于移动和嵌入式设备。

我们可以从 GitHub 下载源代码和预训练模型:
git clone https://github.com/foamliu/MobileFaceNet.git.
cd MobileFaceNet
mkdir weights
cd weights
wget https://github.com/foamliu/MobileFaceNet/releases/download/v1.0/mobilefacenet.pt

为了获得更好的性能和兼容性,可以使用以下示例将 .pt 转换为 .onnx:
convert2onnx.py:

Code Reference

from mobilefacenet import MobileFaceNet
import torch
import time

if __name__ == '__main__':
    filename = 'weights/mobilefacenet.pt'
    print('loading {}...'.format(filename))
    start = time.time()
    model = MobileFaceNet()
    model.load_state_dict(torch.load(filename, map_location=torch.device('cpu')))
    print('elapsed {} sec'.format(time.time() - start))
    print(model)

    output_onnx = 'weights/MobileFaceNet.onnx'
    print("==> Exporting model to ONNX format at '{}'".format(output_onnx))
    input_names = ["input0"]
    output_names = ["output0"]
    inputs = torch.randn(1, 3, 112, 112)

    torch_out = torch.onnx._export(model, inputs, output_onnx, export_params=True, verbose=False,
                                   input_names=input_names, output_names=output_names, opset_version=10)

python3 convert2onnx.py
转换后的 onnx 将位于 weights/MobileFaceNet.onnx。

请先确认已安装QQ通讯软体