expose minimal required versions for ir and opset

Up to now, we set the opset import versions from the values
supported by the onnx library. This is an incorrect behavior.

We now store the maximum opset version required per domain,
based on the operators used in the graph. We then use the
onnx library to compute the minimal required IR version to
execute the model, based on these opset versions.

Fixes #22
Fixes #23

Author: MainRo

ebm2onnx/convert.py

@@ -2,7 +2,6 @@
 import logging
 from enum import Enum
 from copy import deepcopy
-from .utils import get_latest_opset_version
 from ebm2onnx import graph
 from ebm2onnx import ebm
 import ebm2onnx.operators as ops
@@ -114,12 +113,13 @@ def to_graph(model, dtype, name="ebm",
         name: [Optional] The name of the model
         predict_proba: [Optional] For classification models, output prediction probabilities instead of class
         explain: [Optional] Adds an additional output with the score per feature per class
-        target_opset: [Optional] The target onnx opset version to use
+        target_opset: [Optional][Deprecated] The target onnx opset version to use
 
     Returns:
         An ONNX model.
     """
-    target_opset = target_opset or get_latest_opset_version()
+    if target_opset:
+        logging.warning("to_graph: target_opset argument is deprecated")
     root = graph.create_graph(context=context)
 
     inputs = [None for _ in model.feature_names_in_]
@@ -296,23 +296,24 @@ def to_onnx(model, dtype, name="ebm",
         name: [Optional] The name of the model
         predict_proba: [Optional] For classification models, output prediction probabilities instead of class
         explain: [Optional] Adds an additional output with the score per feature per class
-        target_opset: [Optional] The target onnx opset version to use
+        target_opset: [Optional][Deprecated] The target onnx opset version to use
 
     Returns:
         An ONNX model.
     """
+    if target_opset:
+        logging.warning("to_onnx: target_opset argument is deprecated")
     g = to_graph(
         model=model,
         dtype=dtype,
         name=name,
         predict_proba=predict_proba,
         explain=explain,
-        target_opset=target_opset,
         prediction_name=prediction_name,
         probabilities_name=probabilities_name,
         explain_name=explain_name,
         context=context,
         )
 
-    model = graph.to_onnx(g, target_opset, name=name)
+    model = graph.to_onnx(g, name=name)
     return model

ebm2onnx/graph.py

@@ -1,8 +1,8 @@
-from typing import NamedTuple, Callable, Optional, List, Dict, Union
+import logging
+from typing import NamedTuple, Optional, List, Dict, Union
 
 import onnx
-from ebm2onnx import __version__
-from .utils import get_latest_opset_version
+from onnx.helper import make_opsetid
 
 from . import context as _context
 
@@ -14,6 +14,7 @@ class Graph(NamedTuple):
     transients: List[onnx.ValueInfoProto] = []
     nodes: List[onnx.NodeProto] = []
     initializers: List[onnx.TensorProto] = []
+    opsets: Dict[str, int] = {}
 
 
 def extend(i, val):
@@ -52,12 +53,18 @@ def from_onnx(model) -> Graph:
     Returns:
         A Graph object.
     """
+    opsets = {
+        op.domain: op.version
+        for op in model.opset_import
+    }
+
     return Graph(
         context=_context.create(),
         inputs=[n for n in model.graph.input],
         outputs=[n for n in model.graph.output],
         nodes=[n for n in model.graph.node],
         initializers=[n for n in model.graph.initializer],
+        opsets=opsets,
     )
 
 
@@ -74,13 +81,19 @@ def to_onnx(
 
     Args:
         graph: The graph object
-        target_opset: the target opset to use when converting ot onnx, can be an int or a dict
+        target_opset: [Optional][Deprecated] the target opset to use when converting ot onnx, can be an int or a dict
         name: [Optional] An existing ONNX model
 
     Returns:
         A Graph object.
     """
-    #outputs = graph.transients
+    if target_opset:
+        logging.warning("to_onnx: target_opset argument is deprecated")
+
+    opset_imports = [
+        make_opsetid(domain=domain, version=version)
+        for domain,version in graph.opsets.items()
+    ]
 
     graph = onnx.helper.make_graph(
         nodes=graph.nodes,
@@ -89,32 +102,15 @@ def to_onnx(
         outputs=graph.outputs,
         initializer=graph.initializers,
     )
-    model = onnx.helper.make_model(graph, producer_name='ebm2onnx')
-
-    #producer_name = "interpretml/ebm2onnx"
-    #producer_version = __version__
-
-    #domain
-    #model_version
-    #doc_string
 
-    #metadata_props
-
-    # set opset versions
-    if target_opset is not None:
-        if type(target_opset) is int:
-            model.opset_import[0].version = target_opset
-        elif type(target_opset) is dict:
-            del model.opset_import[:]
-
-            for k, v in target_opset.items():
-                opset = model.opset_import.add()
-                opset.domain = k
-                opset.version = v
-        else:
-            raise ValueError(f"ebm2onnx.graph.to_onnx: invalid type for target_opset: {type(target_opset)}.")
-    else:
-        model.opset_import[0].version = get_latest_opset_version()
+    # create the onnx model from the graph.
+    # The onnx library will set the ir version to the minimal required ir that
+    # is compatible with the opset_imports provided.
+    model = onnx.helper.make_model_gen_version(
+        graph,
+        producer_name='ebm2onnx',
+        opset_imports=opset_imports,
+    )
 
     return model
 
@@ -210,4 +206,10 @@ def merge(*args):
                 nodes=extend(g.nodes, graph.nodes),
             )
 
+            # merge opsets, keep higher version for each domain
+            for domain,version in graph.opsets.items():
+                cur_version = g.opsets.get(domain, -1)
+                if version > cur_version:
+                    g.opsets[domain] = version
+
     return g

ebm2onnx/operators.py

@@ -1,8 +1,10 @@
 import onnx
 import ebm2onnx.graph as graph
+from .utils import opset
 
 
 def add():
+    @opset(version=14)
     def _add(g):
         add_result_name = g.context.generate_variable_name('add_result')
         nodes = [
@@ -25,6 +27,7 @@ def _add(g):
 
 
 def argmax(axis=0, keepdims=1, select_last_index=0):
+    @opset(version=13)
     def _argmax(g):        
         argmax_result_name = g.context.generate_variable_name('argmax_result')
         nodes = [
@@ -48,6 +51,7 @@ def _argmax(g):
 
 
 def cast(to):
+    @opset(version=13)
     def _cast(g):
         cast_result_name = g.context.generate_variable_name('cast_result')
         nodes = [
@@ -71,6 +75,7 @@ def _cast(g):
 
 
 def concat(axis):
+    @opset(version=13)
     def _concat(g):
         concat_result_name = g.context.generate_variable_name('concat_result')
 
@@ -96,6 +101,7 @@ def _concat(g):
 
 
 def expand():
+    @opset(version=13)
     def _expand(g):        
         expand_result_name = g.context.generate_variable_name('expand_result')
         nodes = [
@@ -118,6 +124,7 @@ def _expand(g):
 
 
 def flatten(axis=1):
+    @opset(version=13)
     def _flatten(g):        
         flatten_result_name = g.context.generate_variable_name('flatten_result')
         nodes = [
@@ -146,6 +153,7 @@ def gather(axis=0):
         - data
         - indices
     """
+    @opset(version=13)
     def _gather(g):
         gather_result_name = g.context.generate_variable_name('gather_result')
         nodes = [
@@ -169,6 +177,7 @@ def _gather(g):
 
 
 def gather_elements(axis=0):
+    @opset(version=13)
     def _gather_elements(g):
         gather_elements_result_name = g.context.generate_variable_name('gather_elements_result')
         nodes = [
@@ -197,6 +206,7 @@ def gather_nd():
         - scores, as a 2D matrix
         - indices, as a [None, 2] matrix
     """
+    @opset(version=13)
     def _gather_nd(g):
         gather_nd_result_name = g.context.generate_variable_name('gather_nd_result')
         nodes = [
@@ -219,6 +229,7 @@ def _gather_nd(g):
 
 
 def greater_or_equal():
+    @opset(version=16)
     def _greater_or_equal(g):
         greater_or_equal_result_name = g.context.generate_variable_name('greater_or_equal_result')
         nodes = [
@@ -241,6 +252,7 @@ def _greater_or_equal(g):
 
 
 def identity(name, suffix=True):
+    @opset(version=13)
     def _identity(g):
         identity_name = g.context.generate_variable_name(name) if suffix else name
         nodes = [
@@ -263,6 +275,7 @@ def _identity(g):
 
 
 def less():
+    @opset(version=13)
     def _less(g):        
         less_result_name = g.context.generate_variable_name('less_result')
         nodes = [
@@ -285,6 +298,7 @@ def _less(g):
 
 
 def less_or_equal():
+    @opset(version=16)
     def _less_or_equal(g):        
         less_or_equal_result_name = g.context.generate_variable_name('less_or_equal_result')
         nodes = [
@@ -307,6 +321,7 @@ def _less_or_equal(g):
 
 
 def mul():
+    @opset(version=13)
     def _mul(g):        
         mul_result_name = g.context.generate_variable_name('mul_result')
         nodes = [
@@ -329,6 +344,7 @@ def _mul(g):
 
 
 def reduce_sum(keepdims=1, noop_with_empty_axes=0):
+    @opset(version=13)
     def _reduce_sum(g):
         reduce_sum_result_name = g.context.generate_variable_name('reduce_sum_result')
         nodes = [
@@ -354,6 +370,7 @@ def _reduce_sum(g):
 
 
 def reshape(allowzero=0):
+    @opset(version=13)
     def _reshape(g):
         reshape_result_name = g.context.generate_variable_name('reshape_result')
         nodes = [
@@ -376,6 +393,7 @@ def _reshape(g):
 
 
 def softmax(axis=-1):
+    @opset(version=13)
     def _softmax(g):
         softmax_result_name = g.context.generate_variable_name('softmax_result')
         nodes = [
@@ -399,6 +417,7 @@ def _softmax(g):
 
 
 def split(axis=0):
+    @opset(version=18)
     def _split(g):
         split_result_name = [
             g.context.generate_variable_name('split_result')

ebm2onnx/operators_ml.py

@@ -1,8 +1,10 @@
 import onnx
 import ebm2onnx.graph as graph
+from .utils import opset
 
 
 def category_mapper(cats_int64s, cats_strings, default_int64=-1, default_string="_Unused"):
+    @opset(version=1, domain="ai.onnx.ml")
     def _category_mapper(g):
         category_mapper_result_name = g.context.generate_variable_name('category_mapper_result')
         nodes = [

ebm2onnx/sklearn.py

@@ -35,9 +35,10 @@ def convert_ebm_classifier(scope, operator, container):
             context=ctx
         )
 
-        for node in g.nodes:
-            v = container._get_op_version(node.domain, node.op_type)
-            container.node_domain_version_pair_sets.add((node.domain, v))
+        for domain, version in g.opsets.items():
+            container.node_domain_version_pair_sets.add(
+                (domain, version)
+            )
 
         container.nodes.extend(g.nodes)
 

ebm2onnx/utils.py

@@ -1,12 +1,14 @@
-from onnx import defs
-
-
-def get_latest_opset_version():
-    """
-    This module relies on *onnxruntime* to test every
-    converter. The function returns the most recent
-    target opset tested with *onnxruntime* or the opset
-    version specified by *onnx* package if this one is lower
-    (return by `onnx.defs.onnx_opset_version()`).
-    """
-    return min(21, defs.onnx_opset_version())
+
+
+def opset(version: int, domain: str=""):
+    def _operator(op):
+        def _call(g):
+            g = op(g)
+            old_version = g.opsets.get(domain, -1)
+            if version > old_version:
+                g.opsets[domain] = version
+            return g
+
+        return _call
+
+    return _operator