torch2onnx

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)