Skip to content

Commit f4523d2

Browse files
authored
Merge pull request #864 from onnx/gs/ir7
cherry pick fix to set IR based on opset
2 parents f056c74 + 3939afb commit f4523d2

File tree

2 files changed

+11
-0
lines changed

2 files changed

+11
-0
lines changed

tf2onnx/constants.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,3 +37,8 @@
3737

3838
# Environment variables
3939
ENV_TF2ONNX_DEBUG_MODE = "TF2ONNX_DEBUG_MODE"
40+
41+
# Mapping opset to IR version.
42+
OPSET_TO_IR_VERSION = {
43+
1: 3, 2: 3, 3: 3, 4: 3, 5: 3, 6: 3, 7: 4, 8: 4, 9: 4, 10: 5, 11: 6, 12: 7
44+
}

tf2onnx/graph.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -990,6 +990,12 @@ def make_model(self, graph_doc, optimize=False, graph_name="tf2onnx", **kwargs):
990990
kwargs["opset_imports"] = opsets
991991
model_proto = helper.make_model(graph, **kwargs)
992992

993+
# set the IR version based on opset
994+
try:
995+
model_proto.ir_version = constants.OPSET_TO_IR_VERSION.get(self.opset, model_proto.ir_version)
996+
except: # pylint: disable=bare-except
997+
logger.error("ir_version override failed - install the latest onnx version")
998+
993999
# optimize the model proto.
9941000
# TODO: this is disabled by default because of bugs in fuse_consecutive_transposes
9951001
if optimize:

0 commit comments

Comments
 (0)