def wrapped(*args, **kwargs):
global NUMPY_REGISTRATION_SUCCESS
if not NUMPY_REGISTRATION_SUCCESS and mod.has_mod(np, "__version__"):
# We define this alongside load_json/save_json so that it is guaranteed to be
# imported before we need to encode/decode NumPy arrays.
@Encoder.register(np.ndarray)
def encode(array):
outfile = io.BytesIO()
np.save(outfile, array, allow_pickle=False)
outfile.seek(0)
data = base64.b64encode(outfile.read()).decode()
return {"array": data}
@Decoder.register(np.ndarray)
def decode(dct):
def load(mode="base64"):
if mode == "base64":
data = base64.b64decode(dct["array"].encode(), validate=True)
elif mode == "latin-1":
data = dct["array"].encode(mode)
else:
assert False, "Unsupported mode: {:}".format(mode)
infile = io.BytesIO(data)
return np.load(infile, allow_pickle=False)
try:
arr = load()
except:
arr = load("latin-1") # For backwards compatibility
if isinstance(arr, np.ndarray):
return arr
return list(arr.values())[0] # For backwards compatibility
NUMPY_REGISTRATION_SUCCESS = True
return func(*args, **kwargs)
def set_trt_logging_level(sev):
from polygraphy import mod
trt = mod.lazy_import("tensorrt")
if not mod.has_mod(trt, with_attr="__version__"):
return
if sev >= G_LOGGER.CRITICAL:
get_trt_logger().min_severity = trt.Logger.INTERNAL_ERROR
elif sev >= G_LOGGER.ERROR:
get_trt_logger().min_severity = trt.Logger.ERROR
elif sev >= G_LOGGER.INFO:
get_trt_logger().min_severity = trt.Logger.WARNING
elif sev >= G_LOGGER.VERBOSE:
get_trt_logger().min_severity = trt.Logger.INFO
else:
get_trt_logger().min_severity = trt.Logger.VERBOSE
def set_tf_logging_level(sev):
import os
from polygraphy import mod
tf = mod.lazy_import("tensorflow", version="<2.0")
if not mod.has_mod(tf, with_attr="__version__"):
return
if sev > G_LOGGER.WARNING:
tf_sev = tf.compat.v1.logging.ERROR
tf_logging_level = "3"
elif sev > G_LOGGER.INFO:
tf_sev = tf.compat.v1.logging.WARN
tf_logging_level = "2"
elif sev > G_LOGGER.VERBOSE:
tf_sev = tf.compat.v1.logging.INFO
tf_logging_level = "1"
else:
tf_sev = tf.compat.v1.logging.DEBUG
tf_logging_level = "0"
tf.compat.v1.logging.set_verbosity(tf_sev)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = tf_logging_level
def wrapped(*args, **kwargs):
global NUMPY_REGISTRATION_SUCCESS
if not NUMPY_REGISTRATION_SUCCESS and mod.has_mod(np, "__version__"):
# We define this along-side load_json/save_json so that it is guaranteed to be
# imported before we need to encode/decode NumPy arrays.
@Encoder.register(np.ndarray)
def encode(array):
outfile = io.BytesIO()
np.savez(outfile, array)
outfile.seek(0)
return {
"array": outfile.read().decode('latin-1')
}
@Decoder.register(np.ndarray)
def decode(dct):
infile = io.BytesIO(dct["array"].encode('latin-1'))
# We always encode arrays separately.
return list(np.load(infile, allow_pickle=False).values())[0]
NUMPY_REGISTRATION_SUCCESS = True
return func(*args, **kwargs)
def call_impl(self):
"""
Returns:
onnx.ModelProto: The new ONNX model with constants folded.
"""
def run_const_fold_pass(model):
graph = gs_from_onnx(model)
del model
try:
graph.fold_constants(fold_shapes=self.fold_shapes,
partitioning=self.partitioning)
except TypeError as err: # Using an old version of ONNX-GS
if self.partitioning:
G_LOGGER.critical(
"This version of ONNX-GraphSurgeon may not support partitioning the graph.\n"
"Please upgrade to a newer version of ONNX-GraphSurgeon or disable partitioning.\n"
"Note: Error was:\n{:}".format(err))
if self.fold_shapes:
G_LOGGER.critical(
"This version of ONNX-GraphSurgeon may not support folding shapes.\n"
"Please upgrade to a newer version of ONNX-GraphSurgeon or disable shape folding.\n"
"Note: Error was:\n{:}".format(err))
graph.fold_constants()
model = gs.export_onnx(graph.cleanup(), do_type_check=False)
del graph
if self.fold_shapes and self.do_shape_inference:
model = infer_shapes(model)
return model
if not mod.has_mod(onnxrt):
G_LOGGER.error(
"ONNX-Runtime is not installed, so constant folding may be suboptimal or not work at all.\n"
"Consider installing ONNX-Runtime: {:} -m pip install onnxruntime"
.format(sys.executable))
model = self.load()
prefold_num_nodes = len(model.graph.node)
postfold_num_nodes = -1
index = 0
while (prefold_num_nodes != postfold_num_nodes) and (
self.num_passes is None or index < self.num_passes):
prefold_num_nodes = onnx_util.get_num_nodes(model)
G_LOGGER.start("Folding Constants | Pass {:}".format(index + 1))
try:
model = run_const_fold_pass(model)
except Exception as err:
if not self.error_ok:
raise
G_LOGGER.warning(
"Constant folding pass failed. Skipping subsequent passes.\nNote: Error was:\n{:}"
.format(err))
break
else:
postfold_num_nodes = onnx_util.get_num_nodes(model)
index += 1
G_LOGGER.finish(
"\tTotal Nodes | Original: {:5}, After Folding: {:5} | {:5} Nodes Folded"
.format(prefold_num_nodes, postfold_num_nodes,
prefold_num_nodes - postfold_num_nodes))
return model
def np_dtype_from_trt(trt_dtype):
_ = mod.has_mod(np) # Force numpy to be imported
return np.dtype(trt.nptype(trt_dtype))