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。