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通訊軟體