77"""
88import os
99import sys
10- import copy
1110import glob
12- import onnx
13- from onnx import helper
1411
1512
1613def _search_cuda_dir ():
17- paths = os .getenv (' PATH' , '' ).split (os .pathsep )
14+ paths = os .getenv (" PATH" , "" ).split (os .pathsep )
1815 for path in paths :
19- for filename in glob .glob (os .path .join (path , ' cudart64*.dll' )):
16+ for filename in glob .glob (os .path .join (path , " cudart64*.dll" )):
2017 return os .path .dirname (filename )
2118
2219 return None
2320
2421
25- if sys .platform == ' win32' :
22+ if sys .platform == " win32" :
2623 from . import _version # noqa: E402
27- if hasattr (_version , 'cuda' ):
24+
25+ if hasattr (_version , "cuda" ):
2826 cuda_path = _search_cuda_dir ()
2927 if cuda_path is None :
30- raise RuntimeError (
31- "Cannot locate CUDA directory in the environment variable for GPU package" )
28+ raise RuntimeError ("Cannot locate CUDA directory in the environment variable for GPU package" )
3229
3330 os .add_dll_directory (cuda_path )
3431
3532
3633from ._extensions_pydll import ( # noqa
37- PyCustomOpDef , enable_py_op , add_custom_op , hash_64 , default_opset_domain )
34+ PyCustomOpDef ,
35+ enable_py_op ,
36+ add_custom_op ,
37+ hash_64 ,
38+ default_opset_domain ,
39+ )
3840
3941
4042def get_library_path ():
4143 """
4244 The custom operator library binary path
4345 :return: A string of this library path.
4446 """
45- mod = sys .modules [' onnxruntime_extensions._extensions_pydll' ]
47+ mod = sys .modules [" onnxruntime_extensions._extensions_pydll" ]
4648 return mod .__file__
4749
4850
4951class Opdef :
50-
5152 _odlist = {}
5253
5354 def __init__ (self , op_type , func ):
@@ -57,14 +58,14 @@ def __init__(self, op_type, func):
5758
5859 @staticmethod
5960 def declare (* args , ** kwargs ):
60- if len (args ) > 0 and hasattr (args [0 ], ' __call__' ):
61+ if len (args ) > 0 and hasattr (args [0 ], " __call__" ):
6162 raise RuntimeError ("Unexpected arguments {}." .format (args ))
6263 # return Opdef._create(args[0])
6364 return lambda f : Opdef .create (f , * args , ** kwargs )
6465
6566 @staticmethod
6667 def create (func , * args , ** kwargs ):
67- name = kwargs .get (' op_type' , None )
68+ name = kwargs .get (" op_type" , None )
6869 op_type = name or func .__name__
6970 opdef = Opdef (op_type , func )
7071 od_id = id (opdef )
@@ -76,15 +77,15 @@ def create(func, *args, **kwargs):
7677 opdef ._nativedef .op_type = op_type
7778 opdef ._nativedef .obj_id = od_id
7879
79- inputs = kwargs .get (' inputs' , None )
80+ inputs = kwargs .get (" inputs" , None )
8081 if inputs is None :
8182 inputs = [PyCustomOpDef .dt_float ]
8283 opdef ._nativedef .input_types = inputs
83- outputs = kwargs .get (' outputs' , None )
84+ outputs = kwargs .get (" outputs" , None )
8485 if outputs is None :
8586 outputs = [PyCustomOpDef .dt_float ]
8687 opdef ._nativedef .output_types = outputs
87- attrs = kwargs .get (' attrs' , None )
88+ attrs = kwargs .get (" attrs" , None )
8889 if attrs is None :
8990 attrs = {}
9091 elif isinstance (attrs , (list , tuple )):
@@ -106,16 +107,15 @@ def cast_attributes(self, attributes):
106107 elif self ._nativedef .attrs [k ] == PyCustomOpDef .dt_string :
107108 res [k ] = v
108109 else :
109- raise RuntimeError ("Unsupported attribute type {}." .format (
110- self ._nativedef .attrs [k ]))
110+ raise RuntimeError ("Unsupported attribute type {}." .format (self ._nativedef .attrs [k ]))
111111 return res
112112
113113
114114def _on_pyop_invocation (k_id , feed , attributes ):
115115 if k_id not in Opdef ._odlist :
116116 raise RuntimeError (
117- "Unable to find function id={}. "
118- "Did you decorate the operator with @onnx_op?." . format ( k_id ) )
117+ "Unable to find function id={}. " "Did you decorate the operator with @onnx_op?." . format ( k_id )
118+ )
119119 op_ = Opdef ._odlist [k_id ]
120120 rv = op_ .body (* feed , ** op_ .cast_attributes (attributes ))
121121 if isinstance (rv , tuple ):
@@ -127,86 +127,7 @@ def _on_pyop_invocation(k_id, feed, attributes):
127127 res = tuple (res )
128128 else :
129129 res = (rv .shape , rv .flatten ().tolist ())
130- return (k_id , ) + res
131-
132-
133- def _ensure_opset_domain (model ):
134- op_domain_name = default_opset_domain ()
135- domain_missing = True
136- for oi_ in model .opset_import :
137- if oi_ .domain == op_domain_name :
138- domain_missing = False
139-
140- if domain_missing :
141- model .opset_import .extend (
142- [helper .make_operatorsetid (op_domain_name , 1 )])
143-
144- return model
145-
146-
147- def expand_onnx_inputs (model , target_input , extra_nodes , new_inputs ):
148- """
149- Replace the existing inputs of a model with the new inputs, plus some extra nodes
150- :param model: The ONNX model loaded as ModelProto
151- :param target_input: The input name to be replaced
152- :param extra_nodes: The extra nodes to be added
153- :param new_inputs: The new input (type: ValueInfoProto) sequence
154- :return: The ONNX model after modification
155- """
156- graph = model .graph
157- new_inputs = [n for n in graph .input if n .name !=
158- target_input ] + new_inputs
159- new_nodes = list (model .graph .node ) + extra_nodes
160- new_graph = helper .make_graph (
161- new_nodes , graph .name , new_inputs , list (graph .output ), list (graph .initializer ))
162-
163- new_model = copy .deepcopy (model )
164- new_model .graph .CopyFrom (new_graph )
165-
166- return _ensure_opset_domain (new_model )
167-
168-
169- def hook_model_op (model , node_name , hook_func , input_types ):
170- """
171- Add a hook function node in the ONNX Model, which could be used for the model diagnosis.
172- :param model: The ONNX model loaded as ModelProto
173- :param node_name: The node name where the hook will be installed
174- :param hook_func: The hook function, callback on the model inference
175- :param input_types: The input types as a list
176- :return: The ONNX model with the hook installed
177- """
178-
179- # onnx.shape_inference is very unstable, useless.
180- # hkd_model = shape_inference.infer_shapes(model)
181- hkd_model = model
182-
183- n_idx = 0
184- hnode , nnode = (None , None )
185- nodes = list (hkd_model .graph .node )
186- brkpt_name = node_name + '_hkd'
187- optype_name = "op_{}_{}" .format (hook_func .__name__ , node_name )
188- for n_ in nodes :
189- if n_ .name == node_name :
190- input_names = list (n_ .input )
191- brk_output_name = [i_ + '_hkd' for i_ in input_names ]
192- hnode = onnx .helper .make_node (
193- optype_name , n_ .input , brk_output_name , name = brkpt_name , domain = default_opset_domain ())
194- nnode = n_
195- del nnode .input [:]
196- nnode .input .extend (brk_output_name )
197- break
198- n_idx += 1
199-
200- if hnode is None :
201- raise ValueError ("{} is not an operator node name" .format (node_name ))
202-
203- repacked = nodes [:n_idx ] + [hnode , nnode ] + nodes [n_idx + 1 :]
204- del hkd_model .graph .node [:]
205- hkd_model .graph .node .extend (repacked )
206-
207- Opdef .create (hook_func , op_type = optype_name ,
208- inputs = input_types , outputs = input_types )
209- return _ensure_opset_domain (hkd_model )
130+ return (k_id ,) + res
210131
211132
212133PyCustomOpDef .install_hooker (_on_pyop_invocation )
0 commit comments