# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- # This script opens an existing model in onnx format and attempts to # move initializers from model.graph.initializer field to model.graph.sparse_initializer field # and convert them into ONNX COO flat index format. import argparse import logging import numpy as np import sys from typing import Tuple, List import onnx from onnx import ModelProto, SparseTensorProto, TensorProto, numpy_helper logger = logging.getLogger(__name__) real_types = set((int(TensorProto.FLOAT), int(TensorProto.DOUBLE))) def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, type=str, help='input model path') parser.add_argument('--output', required=True, type=str, help='output model path') parser.add_argument('--exclude', required=False, type=str, help='semicolon separated list of initializer names to exclude') parser.add_argument('--tolerance', required=False, type=float, default=1e-6, help='FP absolute tolerance.') parser.add_argument('--sparsity_threshold', required=False, type=float, default=0.5, help='convert to sparse initializers if sparsity is at least this much') parser.add_argument('--verbose', required=False, action='store_true') parser.set_defaults(verbose=False) args = parser.parse_args() return args def setup_logging(verbose): # type: (bool) -> None log_handler = logging.StreamHandler(sys.stdout) if verbose: log_handler.setFormatter(logging.Formatter('[%(filename)s:%(lineno)s - %(funcName)20s()] %(message)s')) logging_level = logging.DEBUG else: log_handler.setFormatter(logging.Formatter('%(filename)20s: %(message)s')) logging_level = logging.INFO log_handler.setLevel(logging_level) logger.addHandler(log_handler) logger.setLevel(logging_level) def convert_tensor_to_sparse(tensor, sparsity_threshold, tolerance): # type: (TensorProto, float, float) -> Tuple[SparseTensorProto, float] """ returns a tuple of sparse_tensor and sparsity level """ values = [] indices = [] nnz_count = 0 tensor_data = numpy_helper.to_array(tensor).flatten() data_len = len(tensor_data) if tensor_data.dtype in real_types: for index in range(data_len): el = tensor_data[index] if abs(el) <= tolerance: values.append(el) indices.append(index) nnz_count += 1 else: for index in range(data_len): el = tensor_data[index] if el != 0: values.append(el) indices.append(index) nnz_count += 1 sparsity = float(1.) - float(nnz_count)/data_len ind_data_type = TensorProto.INT8 ind_dtype = np.int8 ind_len = len(indices) max_indices_value = 0 if ind_len > 0: max_indices_value = indices[-1] if max_indices_value <= np.iinfo(np.int8).max: ind_data_type = TensorProto.INT8 ind_dtype = np.int8 elif max_indices_value <= np.iinfo(np.int16).max: ind_data_type = TensorProto.INT16 ind_dtype = np.int16 elif max_indices_value <= np.iinfo(np.int32).max: ind_data_type = TensorProto.INT32 ind_dtype = np.int32 else: ind_data_type = TensorProto.INT64 ind_dtype = np.int64 logger.debug(f"initializer={tensor.name}, dtype={tensor_data.dtype}, \ data_len={data_len}, nnz={nnz_count}, sparsity={sparsity}, \ max_indices_value={max_indices_value}, sparse_indices_type={ind_dtype}") if sparsity < sparsity_threshold: return (object(), sparsity) tensor_data_bytes = tensor_data.nbytes # create np array and cast data to the appropriate type np_values = np.array(values).astype(tensor_data.dtype) # create np array and cast data to the inferred index type np_indices = np.array(indices).astype(ind_dtype) total_sparse_bytes = np_values.nbytes + np_indices.nbytes logger.debug(f"initializer={tensor.name}, initializer_bytes={tensor_data_bytes}, \ sparse_initializer_bytes={total_sparse_bytes}") # This check is usually useful for sparsity_threshold=0.5 where much # depends on the size of the indices entries and the size of the original tensor. # Big dense tensors command larger indices data type and for large float32 tensors # int32 indices are often selected, thus we really want to guard against loosing # rather than winning. if tensor_data_bytes <= total_sparse_bytes: sparsity = float(1.) - float(tensor_data_bytes)/total_sparse_bytes logger.debug(f"initializer={tensor.name}, adjusted_sparsity={sparsity}") return (object(), sparsity) values_tensor = onnx.helper.make_tensor(tensor.name, tensor.data_type, [len(values)], np_values.tobytes(), raw=True) indicies_tensor = onnx.helper.make_tensor(tensor.name + '_indicies', ind_data_type, [ind_len], np_indices.tobytes(), raw=True) sparse_tensor = onnx.helper.make_sparse_tensor(values_tensor, indicies_tensor, tensor.dims) return (sparse_tensor, sparsity) def convert_initializers(model, exclude_names, sparsity_threshold, tolerance): # type: (ModelProto, List[str], float, float) -> None graph = model.graph converted_sparse = [] remaining_initializers = [] for initializer in graph.initializer: if initializer.name in exclude_names: logger.info(f"initializer={initializer.name} was excluded") continue if initializer.data_type == TensorProto.BOOL: logger.info(f"initializer={initializer.name} contains bool, not converted") remaining_initializers.append(initializer) continue sparse_tensor, sparsity = convert_tensor_to_sparse(initializer, sparsity_threshold, tolerance) if sparsity >= sparsity_threshold: logger.info(f"initializer={initializer.name} converted. sparsity={sparsity}") converted_sparse.append(sparse_tensor) else: remaining_initializers.append(initializer) logger.info(f"initializer={initializer.name} is not converted. sparsity={sparsity}") graph.sparse_initializer.extend(converted_sparse) del graph.initializer[:] graph.initializer.extend(remaining_initializers) def main(): args = parse_arguments() setup_logging(args.verbose) exclude_names = set() if args.exclude is None else set(args.exclude.split(';')) model = ModelProto() with open(args.input, "rb") as input_file: model.ParseFromString(input_file.read()) convert_initializers(model, exclude_names, args.sparsity_threshold, args.tolerance) with open(args.output, "wb") as output_file: s = model.SerializeToString() output_file.write(s) if __name__ == "__main__": main()