| | import numpy as np |
| | import onnx |
| | import torch |
| |
|
| |
|
| | def convert_onnx(net, path_module, output, opset=11, simplify=False): |
| | assert isinstance(net, torch.nn.Module) |
| | img = np.random.randint(0, 255, size=(112, 112, 3), dtype=np.int32) |
| | img = img.astype(np.float) |
| | img = (img / 255. - 0.5) / 0.5 |
| | img = img.transpose((2, 0, 1)) |
| | img = torch.from_numpy(img).unsqueeze(0).float() |
| |
|
| | weight = torch.load(path_module) |
| | net.load_state_dict(weight) |
| | net.eval() |
| | torch.onnx.export(net, img, output, keep_initializers_as_inputs=False, verbose=False, opset_version=opset) |
| | model = onnx.load(output) |
| | graph = model.graph |
| | graph.input[0].type.tensor_type.shape.dim[0].dim_param = 'None' |
| | if simplify: |
| | from onnxsim import simplify |
| | model, check = simplify(model) |
| | assert check, "Simplified ONNX model could not be validated" |
| | onnx.save(model, output) |
| |
|
| | |
| | if __name__ == '__main__': |
| | import os |
| | import argparse |
| | from backbones import get_model |
| |
|
| | parser = argparse.ArgumentParser(description='ArcFace PyTorch to onnx') |
| | parser.add_argument('input', type=str, help='input backbone.pth file or path') |
| | parser.add_argument('--output', type=str, default=None, help='output onnx path') |
| | parser.add_argument('--network', type=str, default=None, help='backbone network') |
| | parser.add_argument('--simplify', type=bool, default=False, help='onnx simplify') |
| | args = parser.parse_args() |
| | input_file = args.input |
| | if os.path.isdir(input_file): |
| | input_file = os.path.join(input_file, "backbone.pth") |
| | assert os.path.exists(input_file) |
| | model_name = os.path.basename(os.path.dirname(input_file)).lower() |
| | params = model_name.split("_") |
| | if len(params) >= 3 and params[1] in ('arcface', 'cosface'): |
| | if args.network is None: |
| | args.network = params[2] |
| | assert args.network is not None |
| | print(args) |
| | backbone_onnx = get_model(args.network, dropout=0) |
| |
|
| | output_path = args.output |
| | if output_path is None: |
| | output_path = os.path.join(os.path.dirname(__file__), 'onnx') |
| | if not os.path.exists(output_path): |
| | os.makedirs(output_path) |
| | assert os.path.isdir(output_path) |
| | output_file = os.path.join(output_path, "%s.onnx" % model_name) |
| | convert_onnx(backbone_onnx, input_file, output_file, simplify=args.simplify) |
| |
|
