import torch
import numpy as np
from nets.retinaface import RetinaFace
from utils.config import cfg_mnet, cfg_re50

model_path = "logs/Epoch150-Total_Loss6.2802.pth"
net = RetinaFace(cfg=cfg_mnet, mode='eval').eval()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net.load_state_dict(torch.load(model_path, map_location=device))

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