import torch
import numpy as np
from nets.facenet import Facenet as facenet

model_path = "logs/ep092-loss0.177-val_loss1.547.pth"
net = facenet(backbone="mobilenet", mode="predict").eval()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net.load_state_dict(torch.load(model_path, map_location=device), strict=False)

img = torch.zeros(1, 3, 160, 160)
torch.onnx.export(net, img, "./facenet.onnx", verbose=False, opset_version=12, input_names=['images'])