Spaces:
Runtime error
Runtime error
| 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 - 0.5) / 0.5 # torch style norm | |
| img = img.transpose((2, 0, 1)) | |
| img = torch.from_numpy(img).unsqueeze(0).float() | |
| weight = torch.load(path_module) | |
| net.load_state_dict(weight, strict=True) | |
| net.eval() | |
| torch.onnx.export( | |
| net, img, output, input_names=["data"], 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, "model.pt") | |
| 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.0, fp16=False, num_features=512) | |
| if args.output is None: | |
| args.output = os.path.join(os.path.dirname(args.input), "model.onnx") | |
| convert_onnx(backbone_onnx, input_file, args.output, simplify=args.simplify) | |