diff --git a/.gitignore b/.gitignore new file mode 100644 index 000000000..1157f9fd0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__ +*.onnx +*.pt +*.pth +*.tar diff --git a/convert_to_onnx.py b/export.py similarity index 55% rename from convert_to_onnx.py rename to export.py index e7f32799a..59c63ef71 100755 --- a/convert_to_onnx.py +++ b/export.py @@ -1,16 +1,8 @@ from __future__ import print_function -import os import argparse import torch -import torch.backends.cudnn as cudnn -import numpy as np from data import cfg_mnet, cfg_re50 -from layers.functions.prior_box import PriorBox -from utils.nms.py_cpu_nms import py_cpu_nms -import cv2 from models.retinaface import RetinaFace -from utils.box_utils import decode, decode_landm -from utils.timer import Timer parser = argparse.ArgumentParser(description='Test') @@ -18,7 +10,9 @@ type=str, help='Trained state_dict file path to open') parser.add_argument('--network', default='mobile0.25', help='Backbone network mobile0.25 or resnet50') parser.add_argument('--long_side', default=640, help='when origin_size is false, long_side is scaled size(320 or 640 for long side)') -parser.add_argument('--cpu', action="store_true", default=True, help='Use cpu inference') +parser.add_argument('--hw', default="cpu", help='device for inference') +parser.add_argument('--mode', default='onnx', help='export to onnx or torch script') +parser.add_argument('output') args = parser.parse_args() @@ -60,6 +54,7 @@ def load_model(model, pretrained_path, load_to_cpu): if __name__ == '__main__': + args = parser.parse_args() torch.set_grad_enabled(False) cfg = None if args.network == "mobile0.25": @@ -67,22 +62,44 @@ def load_model(model, pretrained_path, load_to_cpu): elif args.network == "resnet50": cfg = cfg_re50 # net and model - net = RetinaFace(cfg=cfg, phase = 'test') - net = load_model(net, args.trained_model, args.cpu) + net = RetinaFace(cfg=cfg, phase='test') + load_to_cpu = args.hw == 'cpu' or args.hw == 'eia' + net = load_model(net, args.trained_model, load_to_cpu) net.eval() print('Finished loading model!') print(net) - device = torch.device("cpu" if args.cpu else "cuda") - net = net.to(device) + if args.mode == 'onnx': + device = torch.device("cpu" if args.hw == 'cpu' else "cuda") + net = net.to(device) # ------------------------ export ----------------------------- - output_onnx = 'FaceDetector.onnx' - print("==> Exporting model to ONNX format at '{}'".format(output_onnx)) - input_names = ["input0"] - output_names = ["output0"] - inputs = torch.randn(1, 3, args.long_side, args.long_side).to(device) - - torch_out = torch.onnx._export(net, inputs, output_onnx, export_params=True, verbose=False, - input_names=input_names, output_names=output_names) - + if args.mode == 'onnx': + output_onnx = args.output + print("==> Exporting model to ONNX format at '{}'".format(output_onnx)) + input_names = ["input0"] + output_names = ["output0"] + inputs = torch.randn(1, 3, args.long_side, args.long_side).to(device) + torch_out = torch.onnx._export(net, inputs, output_onnx, + export_params=True, verbose=False, + input_names=input_names, + output_names=output_names, + opset_version=11) + elif args.mode == 'tscript': + output_mod = args.output + # scripted_model = torch.jit.script(net) + # torch.jit.save(scripted_model, output_mod) + inputs = torch.randn(1, 3, args.long_side, args.long_side) + check_inputs = [torch.randn(1, 3, args.long_side, args.long_side), + torch.randn(2, 3, args.long_side, args.long_side)] + if args.hw == 'cpu' or args.hw == 'gpu': + device = torch.device("cpu" if args.hw == 'cpu' else "cuda") + net = net.to(device) + traced_model = torch.jit.trace(net, inputs, + check_inputs=check_inputs) + elif args.hw == 'eia': + with torch.jit.optimized_execution(True, + {'target_device': 'eia:0'}): + traced_model = torch.jit.trace(net, inputs, + check_inputs=check_inputs) + torch.jit.save(traced_model, output_mod) diff --git a/export_models.sh b/export_models.sh new file mode 100755 index 000000000..abaead2b5 --- /dev/null +++ b/export_models.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +python3 export.py -m weights/mobilenet0.25_Final.pth --network mobile0.25 --mode tscript mnet.0.25.pt +python3 export.py -m weights/Resnet50_Final.pth --network resnet50 --mode tscript resnet50.pt + +python3 export.py -m weights/mobilenet0.25_Final.pth --network mobile0.25 mnet.0.25.onnx +python3 export.py -m weights/Resnet50_Final.pth --network resnet50 resnet50.onnx diff --git a/weights/download_all_pretrain.sh b/weights/download_all_pretrain.sh new file mode 100755 index 000000000..9a789f70d --- /dev/null +++ b/weights/download_all_pretrain.sh @@ -0,0 +1,4 @@ +#!/bin/bash +./google_download.sh 14KX6VqF69MdSPk3Tr9PlDYbq7ArpdNUW Resnet50_Final.pth +./google_download.sh 1q36RaTZnpHVl4vRuNypoEMVWiiwCqhuD mobilenetV1X0.25_pretrain.tar +./google_download.sh 15zP8BP-5IvWXWZoYTNdvUJUiBqZ1hxu1 mobilenet0.25_Final.pth diff --git a/weights/google_download.sh b/weights/google_download.sh new file mode 100755 index 000000000..76b28ed98 --- /dev/null +++ b/weights/google_download.sh @@ -0,0 +1,8 @@ +#!/bin/bash +set -euox pipefail +FILEID=$1 +FILENAME=$2 +cookie=$(FILEID=$1 ; wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies \ + --no-check-certificate "https://docs.google.com/uc?export=download&id=$FILEID" -O- | \ + sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p') +wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$cookie&id=$FILEID" -O $FILENAME && rm -rf /tmp/cookies.txt