def test_nested(self): self.assertTrue(nest.is_nested(list())) self.assertTrue(nest.is_sequence(list())) self.assertTrue(nest.is_nested(tuple())) self.assertTrue(nest.is_sequence(tuple())) self.assertTrue(nest.is_nested(dict())) self.assertFalse(nest.is_sequence(dict()))
def assertEqual( self, first, second, msg=None, prec=None, ): if prec is None: prec = self.precision inputs = nest.flatten(first) num_first = len(inputs) inputs += nest.flatten(second) num_second = len(inputs) - num_first for i, input in enumerate(inputs): if isinstance(input, torch.Tensor): inputs[i] = input.numpy() first = inputs[:num_first] if num_first > 1 else inputs[0] second = inputs[num_first:len(inputs)] if num_second > 1 else inputs[ num_first] if isinstance(first, np.ndarray) and isinstance(second, np.ndarray): super(OpTestCase, self).assertEqual(first.shape, second.shape) if first.dtype == np.bool and second.dtype == np.bool: diff = first ^ second num_unique = len(np.unique(diff)) self.assertLessEqual(num_unique, 1, msg) else: diff = np.abs(first - second) max_err = diff.max() self.assertLessEqual(max_err, prec, msg) elif nest.is_sequence(first) and nest.is_sequence(second): for a, b in zip(first, second): self.assertEqual(a, b, msg, prec) else: super(OpTestCase, self).assertEqual(first, second, msg)
def normalize_paddings(value, rank): """Repeat the paddings according to the rank.""" if isinstance(value, int): return ((value, value), ) * rank elif nest.is_sequence(value): if isinstance(value[0], int): return normalize_tuple(value, rank) elif nest.is_sequence(value[0]): value = [normalize_tuple(v, 2) for v in value] if len(value) > rank: return (value[i] for i in range(rank)) else: return tuple([value[i] for i in range(len(value))] + [value[-1] for _ in range(len(value), rank)])
def run(self, inputs, **kwargs): """Run the model. Parameters ---------- inputs : Union[Sequence, Dict] The input arrays. Returns ------- namedtuple The model outputs. """ if isinstance(inputs, numpy.ndarray): inputs = [inputs] if isinstance(inputs, dict): for name, value in inputs.items(): self._input_dict[name]._impl.FromNumpy(value) elif nest.is_sequence(inputs): for ref, value in zip(self._input_dict.values(), inputs): ref._impl.FromNumpy(value) else: raise ValueError('Excepted sequence or dict inputs.') self._function.callback(return_outputs=False) named_outputs = namedtupledict('Outputs', list(self._output_dict.keys())) return named_outputs(*(self._output_dict.values()))
def run(self, inputs, **kwargs): """Run the model. Parameters ---------- inputs : Union[Sequence, Dict] The input arrays. Returns ------- namedtuple The model outputs. """ if isinstance(inputs, numpy.ndarray): inputs = [inputs] if isinstance(inputs, dict): for name, value in inputs.items(): self._input_dict[name]._impl.FromNumpy(value) elif nest.is_sequence(inputs): for ref, value in zip(self._input_dict.values(), inputs): ref._impl.FromNumpy(value) else: raise ValueError('Excepted sequence or dict inputs.') self._context.run() return self._output_tuple(*self._output_dict.values())
def split(tensor, split_size_or_sections, dim=0, copy=True): """Split input into chunks along the given dimension. Either size of every chunk or each chunk will be accepted: ```python x = torch.tensor([1, 2, 3, 4, 5, 6]) # Shape: (6,) -> (4,), (2,) print(torch.split(x, split_size_or_sections=4)) # Shape: (6,) -> (5,), (1,) print(torch.split(x, split_size_or_sections=(5, 1))) ``` :attr:`dim` can be negative: ```python x = torch.tensor([[1, 2, 3], [4, 5, 6]]) print(torch.split(x, 2, dim=1)) print(torch.split(x, 2, dim=-1)) # Equivalent ``` Parameters ---------- tensor : dragon.vm.torch.Tensor The input tensor. split_size_or_sections : Union[int, Sequence[int] The number or size of chunks. dim : int, optional, default=0 The dimension to split. copy : bool, optional, default=True Copy or create the views of input. Returns ------- Sequence[dragon.vm.torch.Tensor] The output tensors. """ if nest.is_sequence(split_size_or_sections): size_splits = split_size_or_sections num_splits = len(split_size_or_sections) else: size = tensor.shape[dim] if size % split_size_or_sections == 0: num_splits = size // split_size_or_sections size_splits = [split_size_or_sections] * num_splits else: num_splits = size // split_size_or_sections + 1 size_splits = [split_size_or_sections] * num_splits size_splits[-1] = size - (split_size_or_sections * (num_splits - 1)) return Function.apply('Split', tensor.device, [tensor], outputs=[None] * num_splits, axis=dim, num_splits=num_splits, split=size_splits, copy=copy)
def split(tensor, split_size_or_sections, dim=0): """Split input into chunks along the given dimension. Either size of every chunk or each chunk will be accepted: ```python x = torch.tensor([1, 2, 3, 4, 5, 6]) # Shape: (6,) -> (4,), (2,) print(torch.split(x, split_size_or_sections=4)) # Shape: (6,) -> (5,), (1,) print(torch.split(x, split_size_or_sections=(5, 1))) ``` The ``dim`` can be negative representing the last-k axis: ```python x = torch.tensor([[1, 2], [3, 4], [5, 6]]) print(torch.split(x, 2, dim=1)) print(torch.split(x, 2, dim=-1)) # Equivalent ``` Parameters ---------- tensor : dragon.vm.torch.Tensor The input tensor. split_size_or_sections : Union[int, Sequence[int] The number or size of chunks. dim : int, optional, default=0 The dimension to split. Returns ------- Sequence[dragon.vm.torch.Tensor] The output tensors. """ if nest.is_sequence(split_size_or_sections): num_splits = len(split_size_or_sections) size_splits = split_size_or_sections else: size = tensor.shape[dim] if size % split_size_or_sections == 0: num_splits = size // split_size_or_sections size_splits = [split_size_or_sections] * num_splits else: num_splits = size // split_size_or_sections + 1 size_splits = [split_size_or_sections] * num_splits size_splits[-1] = size - (split_size_or_sections * (num_splits - 1)) return _functions.Split \ .instantiate( tensor.device, axis=dim, size_splits=size_splits, ).apply(tensor, num_splits)
def _maybe_build(self, inputs): if not self.built: input_spec.assert_input_compatibility(self.input_spec, inputs, self.name) inputs_list = nest.flatten(inputs) input_shapes = None if all(hasattr(x, 'shape') for x in inputs_list): input_shapes = [x.shape for x in inputs_list] if not nest.is_sequence(inputs): input_shapes = input_shapes[0] self.build(input_shapes)
def split(inputs, num_or_size_splits, axis=0, copy=True, **kwargs): """Split input into chunks along the given axis. Either number or size of splits will be accepted: ```python x = dragon.constant([[1, 2], [3, 4], [5, 6]]) # Shape: (3, 2) -> (2, 2), (1, 2) print(dragon.split(x, num_or_size_splits=2)) # Shape: (3, 2) -> (1, 2), (2, 2) print(dragon.split(x, num_or_size_splits=(1, 2))) ``` :attr:`axis` can be negative: ```python x = dragon.constant([[1, 2], [3, 4], [5, 6]]) print(dragon.split(x, 2, axis=1)) print(dragon.split(x, 2, axis=-1)) # Equivalent ``` Parameters ---------- inputs : dragon.Tensor The input tensor. num_or_size_splits: Union[int, Sequence[int]] The number or size of chunks. axis : int, optional, default=0 The axis to split. copy : bool, optional, default=True Copy or create the views of input. Returns ------- Sequence[dragon.Tensor] The output tensors. """ if nest.is_sequence(num_or_size_splits): size_splits = num_or_size_splits num = num_splits = len(num_or_size_splits) else: size_splits = None num, num_splits = num_or_size_splits, 0 if context.executing_eagerly(): return OpLib.execute( 'Split', inputs, outputs=[None] * num, axis=axis, num_splits=num_splits, split=size_splits, copy=copy) return OpLib.add('Split', inputs, num_outputs=num, axis=axis, split=size_splits, copy=copy, **kwargs)
def shape(self, value): """Set the tensor shape. Parameters --------- value : Sequence[int] The shape to set. """ if value is not None: if not nest.is_sequence(value): raise TypeError( 'The <shape> should be a sequence. Got {}.'.format( type(value).__name__)) self._shape = tuple(nest.flatten(value)) else: self._shape = value
def __call__(self, inputs, **kwargs): """The preprocessor for ``self.forward(...)``.""" # Maybe build the layer at the first time. if not self._built: input_list = nest.flatten(inputs) input_shapes = None if all(hasattr(x, 'shape') for x in input_list): input_shapes = [x.shape for x in input_list] if not nest.is_sequence(inputs): input_shapes = input_shapes[0] self.build(input_shapes) # Call the forward implementation to get outputs. outputs = self.forward(inputs, **kwargs) # Record the nodes if necessary. if not self._nodes_fixed: self._add_node(inputs, outputs) return outputs
def graph_def_to_onnx_graph( cls, graph_def, input_names=None, output_names=None, input_shapes=None, constants=None, value_info=None, opset_version=None, workspace=None, verbose=True, ): input_names = [] if input_names is None else input_names output_names = [] if output_names is None else output_names constants = {} if constants is None else constants value_info = {} if value_info is None else value_info if not nest.is_sequence(input_names): raise ValueError('<input_names> should be a sequence.') if not nest.is_sequence(output_names): raise ValueError('<output_names> should be a sequence.') if not isinstance(constants, dict): raise ValueError('<constants> should be a dict with name -> value.') if not isinstance(value_info, dict): raise ValueError('<value_info> should be a dict with name -> (dtype, shape).') # Determine the opset version to select exporters. if opset_version is None: opset_version = cls._check_opset_version(opset_version) # Create aliases for blobs. blob_aliases = {} for i, alias in enumerate(output_names): blob_aliases[graph_def.output[i]] = alias workspace.RegisterAlias(graph_def.output[i], alias) if graph_def.output[i] in value_info: value_info[alias] = value_info[graph_def.output[i]] for i, alias in enumerate(input_names): blob_aliases[graph_def.input[i]] = alias workspace.RegisterAlias(graph_def.input[i], alias) if graph_def.input[i] in value_info: value_info[alias] = value_info[graph_def.input[i]] # Maybe rewrite the input shapes for future development. # A common case is that we should fill ``-1`` for dynamic dimension # in the inference runtime like TensorRT. if input_shapes is not None: if isinstance(input_shapes, dict): for k, v in input_shapes.items(): value_info[k] = (value_info[k][0], v) else: for k, v in zip(graph_def.input[:], input_shapes): value_info[k] = (value_info[k][0], v) # Prepare to make the graph. onnx_graph = onnx.GraphProto(name=graph_def.name if len(graph_def.name) > 0 else 'onnx-model') blob_shapes, blob_names = {}, {} blob_versions = collections.defaultdict( int, **dict((blob_aliases.get(k, k), 1) for k in helper.collect_inputs(graph_def))) initializers, seen_initializers = [], set() # Build translator context. context = export_util.TranslatorContext( workspace=workspace, blob_names=blob_names, blob_shapes=blob_shapes, blob_versions=blob_versions, opset_version=opset_version, ) # Add nodes. for op in graph_def.op: # Get the shape of inputs and outputs. for name in itertools.chain(op.input, op.output): impl = workspace.GetTensor(name) if impl is not None: blob_shapes[name] = impl.dims else: blob_shapes[name] = value_info[name][1] # Translate definition. nodes, const_tensors = cls._make_node(op, context) # Rewritten for names. for node in nodes: node.input[:] = [blob_aliases.get(e, e) for e in node.input] node.output[:] = [blob_aliases.get(e, e) for e in node.output] cls._rewrite_for_ssa(node, context) # Convert constant outputs if necessary. if None in nodes: const_tensors = [helper.from_tensor(name, workspace) for name in op.output] else: onnx_graph.node.extend(nodes) # Merge constant tensors. if const_tensors is not None: value_info = {**value_info, **dict((e.name, (e.data_type, e.dims)) for e in const_tensors)} for tensor in const_tensors: if tensor.name not in seen_initializers: initializers.append(tensor) seen_initializers.add(tensor.name) # Add constants. if constants is not None: for k, v in constants.items(): initializers.append(helper.from_array(v, name=k)) # Add inputs. for name in helper.collect_inputs(onnx_graph): try: onnx_graph.input.extend([ helper.make_tensor_value_info( name=name, elem_type=value_info[name][0], shape=value_info[name][1])]) except KeyError: impl = workspace.GetTensor(name) if impl is not None: initializer = helper.from_tensor(name, workspace) onnx_graph.input.extend([ helper.make_tensor_value_info( name=name, elem_type=initializer.data_type, shape=initializer.dims)]) if name not in seen_initializers: initializers.append(initializer) seen_initializers.add(initializer.name) else: raise ValueError( 'Info of tensor `{}` is missing, ' 'specify it in <value_info>.'.format(name)) # Add initializers. onnx_graph.initializer.extend(initializers) # Add outputs. onnx_graph.output.extend( helper.make_tensor_value_info( name=blob_names.get(name_v2, name_v2), elem_type=value_info[name_v2][0], shape=value_info[name_v2][1]) for name_v2 in [blob_aliases.get(name, name) for name in set(graph_def.output)]) if verbose: print(helper.printable_graph(onnx_graph)) return onnx_graph
def checkpoint_sequential(functions, input, segments=1, **kwargs): """Apply functions and create segmental checkpoints. Parameters ---------- functions : Union[torch.nn.Sequential, Sequence[callable]] The functions to apply sequentially. input : dragon.vm.torch.Tensor The input tensor. segments : Union[int, Sequence[int]], optional The number or size of chunked checkpoints. Returns ------- Any The function outputs. """ def run_function(start, end, functions): def forward(input): for j in range(start, end): input = functions[j](input) with no_checkpoint(): input = functions[end](input) return input return forward preserve_rng_state = kwargs.pop('preserve_rng_state', True) variable_scope = kwargs.pop('variable_scope', 'Buffer') if kwargs: raise ValueError('Unexpected keyword arguments: ' + ','.join(arg for arg in kwargs)) if isinstance(functions, Sequential): functions = list(functions.children()) start, end = 0, len(functions) - 1 if not grad_mode.is_grad_enabled(): return run_function(start, end, functions)(input) if nest.is_sequence(segments): size_segments = segments if sum(size_segments) != len(functions): raise ValueError('Failed to chunk {} functions into {} segments.' .format(len(functions), segments)) else: size = (len(functions) + segments - 1) // segments last_size = len(functions) - size * (segments - 1) if last_size <= 0: raise ValueError('Failed to chunk {} functions into {} segments.' .format(len(functions), segments)) size_segments = [size] * (segments - 1) + [last_size] for size in size_segments: end = start + size - 1 input = checkpoint( run_function(start, end, functions), input, preserve_rng_state=preserve_rng_state, variable_scope=variable_scope) start = end + 1 return input