import json import os import typing from collections import OrderedDict import numpy as np import torch import torch.nn as nn from tinynn.util.util import get_logger log = get_logger(__name__) def tensor_config(tensors: typing.List[torch.Tensor], transpose: typing.List[bool], with_shape: bool) -> OrderedDict: """Generate the tensor info needed for the config file for the TinyNeuralNetwork converter Args: tensors (typing.List[torch.Tensor]): The tensors to gather info transpose (typing.List[bool]): Whether to perform nchw-nhwc transpose for the tensors with_shape (bool): Whether to dump the shape of the tensor Returns: OrderedDict: The info of the tensors """ tensor_list = [] for t, trans in zip(tensors, transpose): tensor_dict = OrderedDict() if with_shape: tensor_dict['shape'] = list(t.shape) try: tensor_dict['type'] = str(t.detach().numpy().dtype) except Exception: type_str = str(t.dtype) if 'torch.q' in type_str: type_str = type_str.replace('torch.q', '') else: type_str = type_str.replace('torch.', '') tensor_dict['type'] = type_str if trans is None: trans = len(t.shape) == 4 tensor_dict['transpose'] = trans tensor_list.append(tensor_dict) return tensor_list def generate_converter_config( inputs: typing.List[torch.Tensor], outputs: typing.List[torch.Tensor], input_transpose: typing.Union[typing.Iterable[bool], bool], output_transpose: typing.Union[typing.Iterable[bool], bool], export_file: str, tflite_path: typing.Optional[str] = None, config_path: typing.Optional[str] = None, ): """ Generate a config file that will work for the TinyNeuralNetwork converter Args: inputs (typing.List[torch.Tensor]): The input tensors outputs (typing.List[torch.Tensor]): The output tensors input_transpose (typing.Union[typing.Iterable[bool], bool]): The flag whether to insert transpose after the \ input nodes output_transpose (typing.Union[typing.Iterable[bool], bool]): The flag whether to insert transpose after the \ output nodes export_file (str): The path of the generate torchscript model tflite_path (str): The path of the generate tflite model. Defaults to None. config_path (str): The path of the generate config. Defaults to None. Raises: AssertionError: input transpose should either be boolean or list of booleans, under the latter condition, \ the size of the list should be the same of that of the inputs """ if type(input_transpose) in (tuple, list): if len(input_transpose) != len(inputs) or not all((type(x) is bool for x in input_transpose)): raise AssertionError('input transpose should either be boolean or list of booleans') elif type(input_transpose) is bool or input_transpose is None: input_transpose = [input_transpose] * len(inputs) else: raise AssertionError('input transpose should either be boolean or list of booleans') if type(output_transpose) in (tuple, list): if len(output_transpose) != len(outputs) or not all((type(x) is bool for x in output_transpose)): raise AssertionError('output transpose should either be boolean or list of booleans') elif type(output_transpose) is bool or output_transpose is None: output_transpose = [output_transpose] * len(inputs) else: raise AssertionError('output transpose should either be boolean or list of booleans') if tflite_path is None: tflite_path = export_file.replace('.pt', '.tflite') json_obj = OrderedDict( { 'src_model': export_file, 'dst_model': tflite_path, 'inputs': tensor_config(inputs, input_transpose, True), 'outputs': tensor_config(outputs, output_transpose, False), } ) if config_path is None: config_path = export_file.replace('.pt', '.json') with open(config_path, 'w') as f: json.dump(json_obj, f, indent=4) def export_converter_files( model: nn.Module, dummy_input: typing.Union[torch.Tensor, typing.Iterable[torch.Tensor]], export_dir: typing.Optional[str] = None, model_name: typing.Optional[str] = None, input_transpose: typing.Optional[typing.Union[bool, typing.Iterable[bool]]] = None, output_transpose: typing.Optional[typing.Union[bool, typing.Iterable[bool]]] = None, dump_graph: bool = False, ): """ Automatically generate required files for the model converter Args: model (nn.Module): The input model dummy_input (typing.Union[torch.Tensor, typing.Iterable[torch.Tensor]]): A viable input to the model export_dir (typing.Optional[str], optional): Directory to use for exporting. Defaults to None(os.getcwd()). model_name (typing.Optional[str], optional): File name for exporting. Defaults to None("jit_model"). input_transpose (typing.Optional[typing.Union[bool, typing.Iterable[bool]]], optional): Whether to transpose \ the input(s). Defaults to None(True for 4d-input, False otherwise). output_transpose (typing.Optional[typing.Union[bool, typing.Iterable[bool]]], optional): Whether to transpose \ the input(s). Defaults to None(True for 4d-output, False otherwise). dump_graph (bool, optional): Whether to print the traced graph. Defaults to False. """ if export_dir is None: export_dir = os.getcwd() if model_name is None: model_name = 'jit_model' export_file = os.path.abspath(os.path.join(export_dir, f'{model_name}.pt')) script = torch.jit.trace(model, dummy_input) if dump_graph: log.info(script.inlined_graph) os.makedirs(export_dir, exist_ok=True) torch.jit.save(script, export_file) if type(dummy_input) not in (tuple, list): inputs = [dummy_input] else: inputs = dummy_input output = model(*inputs) if type(output) not in (tuple, list): outputs = [output] else: new_output = [] for item in output: if type(item) in (tuple, list): new_output.extend(item) else: new_output.append(item) outputs = new_output generate_converter_config(inputs, outputs, input_transpose, output_transpose, export_file) def get_tensor_details(config): input_shapes = [] input_transpose = [] input_dtypes = [] output_transpose = [] input_shapes.extend((inp['shape'] for inp in config['inputs'])) input_transpose.extend((inp['transpose'] for inp in config['inputs'])) input_dtypes.extend((inp['type'] for inp in config['inputs'])) output_transpose.extend((outp['transpose'] for outp in config['outputs'])) return input_shapes, input_transpose, input_dtypes, output_transpose def prepare_input_arrays(input_shapes, input_transpose, input_dtypes): inputs = [] for i in range(len(input_shapes)): input_shape = input_shapes[i] tranpose = input_transpose[i] input_data = np.zeros(input_shape, dtype=input_dtypes[i]) if tranpose: input_data = input_data.transpose((0, 2, 3, 1)) inputs.append(input_data) return inputs def data_to_pytorch(inputs, input_transpose): torch_inputs = list(map(torch.from_numpy, inputs)) for i in range(len(torch_inputs)): if input_transpose[i]: torch_inputs[i] = torch_inputs[i].permute(0, 3, 1, 2) return torch_inputs def parse_config( json_file: str, prepare_inputs: bool = True ) -> typing.Tuple[str, str, typing.List[bool], typing.List[torch.Tensor], typing.List[bool]]: """Parses the configuration file for converter Args: json_file (str): The path of the configuration file prepare_inputs (bool, optional): Whether to prepare inputs. Defaults to True. Returns: tflite_model_path (str): The path used to place the generated tflite model torch_model_path (str): The path of the torchscript model input_transpose (typing.List[bool]): Flag variables whether the inputs will be transposed (nchw -> nhwc) torch_inputs (typing.List[torch.Tensor]): The prepared inputs if prepare_inputs is True, otherwise None output_transpose (typing.List[bool]): Flag variables whether the outputs will be transposed (nchw -> nhwc) """ with open(json_file, 'r') as f: config = json.load(f) tflite_model_path = config['dst_model'] torch_model_path = config['src_model'] if prepare_inputs: input_shapes, input_transpose, input_dtypes, output_transpose = get_tensor_details(config) inputs = prepare_input_arrays(input_shapes, input_transpose, input_dtypes) torch_inputs = data_to_pytorch(inputs, input_transpose) else: torch_inputs = None return torch_model_path, tflite_model_path, input_transpose, torch_inputs, output_transpose