1
0
Fork 0

feat(api): add optimization flag to ORT extended optimizations

This commit is contained in:
Sean Sube 2023-12-26 20:22:32 -06:00
parent 006f983269
commit e8a30ae56b
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 3 additions and 0 deletions

View File

@ -161,6 +161,9 @@ class DeviceParams:
elif "onnx-graph-basic" in self.optimizations: elif "onnx-graph-basic" in self.optimizations:
logger.debug("enabling basic ONNX graph optimizations") logger.debug("enabling basic ONNX graph optimizations")
sess.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC sess.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC
elif "onnx-graph-extended" in self.optimizations:
logger.debug("enabling extended ONNX graph optimizations")
sess.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_EXTENDED
elif "onnx-graph-all" in self.optimizations: elif "onnx-graph-all" in self.optimizations:
logger.debug("enabling all ONNX graph optimizations") logger.debug("enabling all ONNX graph optimizations")
sess.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL sess.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL