import numpy as np import onnx import torch def change_input_dim(model): # Use some symbolic name not used for any other dimension sym_batch_dim = "N" # or an actal value actual_batch_dim = 4 # The following code changes the first dimension of every input to be batch-dim # Modify as appropriate ... note that this requires all inputs to # have the same batch_dim inputs = model.graph.input for input in inputs: # Checks omitted.This assumes that all inputs are tensors and have a shape with first dim. # Add checks as needed. dim1 = input.type.tensor_type.shape.dim[0] # update dim to be a symbolic value dim1.dim_param = sym_batch_dim # or update it to be an actual value: # dim1.dim_value = actual_batch_dim def apply(transform, infile, outfile): model = onnx.load(infile) transform(model) onnx.save(model, outfile) #apply(change_input_dim, r"input-file-name, r"output-file-name") # used to change the batch_size
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 # torch style norm img = img.transpose((2, 0, 1)) img = torch.from_numpy(img).unsqueeze(0).float() input_h, input_w = 112, 112
dummy_input = torch.randn(1, 3, 112, 112) weight = torch.load(path_module) net.load_state_dict(weight) net.eval() torch.onnx.export(net, img, output, input_names=['input_batch'], 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' ### change the batch_size graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1 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 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) ## TODO need to load your model output_path = args.outputif 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)