File size: 730 Bytes
b5fdc16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import time
import onnx
from onnxscript import optimizer
def timeit(f, message):
def timed(*args, **kw):
ts = time.time()
result = f(*args, **kw)
te = time.time()
print(f"{message} time: {te - ts}")
return result
return timed
load = timeit(onnx.load, "Load")
save = timeit(onnx.save, "Save")
infer = timeit(onnx.shape_inference.infer_shapes, "Infer")
fold_constants = timeit(optimizer.fold_constants, "Fold Constants")
remove_unused = timeit(optimizer.remove_unused_nodes, "Remove Unused")
optimize = timeit(optimizer.optimize, "Optimize")
# rewrite = timeit(all_rules.apply_to_model, "Rewrite")
|