codekingpro/portable-devtools

1import argparse
2
3import onnx
4
5
6def get_args():
7    parser = argparse.ArgumentParser()
8    parser.add_argument("--input", required=True, help="input model")
9    parser.add_argument("--output", required=True, help="output model")
10    args = parser.parse_args()
11    return args
12
13
14def remove_initializer_from_input(model: onnx.ModelProto) -> bool:
15    if model.ir_version < 4:
16        print("Model with ir_version below 4 requires to include initializer in graph input")
17        return False
18
19    inputs = model.graph.input
20    name_to_input = {}
21    for input in inputs:
22        name_to_input[input.name] = input
23
24    modified = False
25    for initializer in model.graph.initializer:
26        if initializer.name in name_to_input:
27            modified = True
28            inputs.remove(name_to_input[initializer.name])
29
30    return modified
31
32
33if __name__ == "__main__":
34    args = get_args()
35    model = onnx.load(args.input)
36    remove_initializer_from_input(model)
37    onnx.save(model, args.output)
38