def reshape_shape(op): # type: (Caffe2Operation)->ShapeResult if len(op.inputs) == 1: shape = op.attribs['shape'] elif len(op.inputs) == 2: if op.inputs[1].data is None: raise utils.NNEFToolsException( 'Reshape is not supported with calculated shape.') shape = op.inputs[1].data.tolist() else: assert False graph_utils.replace_tensor_in_consumers( op.graph, op.outputs[1], Caffe2Tensor(graph=op.graph, shape=[op.inputs[0].rank], data=np.array(op.inputs[0].shape, dtype=np.int64), dtype=DTYPE_INT64), remove=False) op.attribs['shape'] = shape op.inputs = (op.inputs[0], ) return (infer.reshape(op.inputs[0].shape, shape=shape, zero_means_same=True), [op.inputs[0].rank]), \ (op.inputs[0].dtype, DTYPE_INT64)
def remove_passthrough_ex(g, op): # type: (BaseGraph, BaseOperation)->None op_input = op.inputs[0] op_output = op.outputs[0] g.remove_operation(op, unlink=True) graph_utils.replace_tensor_in_consumers(g, op_output, op_input, remove=True)
def shape_shape(op): # type: (Caffe2Operation)->ShapeResult graph_utils.replace_tensor_in_consumers(op.graph, op.output, Caffe2Tensor(graph=op.graph, shape=[op.input.rank], data=np.array( op.input.shape, dtype=np.int64), dtype=DTYPE_INT64), remove=False) return [op.input.rank], DTYPE_INT64
def size_shape(op): # type: (Caffe2Operation)->ShapeResult graph_utils.replace_tensor_in_consumers( op.graph, op.outputs[0], Caffe2Tensor(graph=op.graph, shape=[], data=np.array(op.inputs[0].count, dtype=np.int64), dtype=DTYPE_INT64), remove=False) return one_element_0d_shape(op)
def concat_shape(op): # type: (Caffe2Operation)->ShapeResult if op.attribs['add_axis']: output_shape = infer.stack([input.shape for input in op.inputs], axis=op.attribs['axis']) else: output_shape = infer.concat([input.shape for input in op.inputs], axis=op.attribs['axis']) graph_utils.replace_tensor_in_consumers( op.graph, op.outputs[1], Caffe2Tensor( graph=op.graph, shape=[len(op.inputs)], data=np.array( [input.shape[op.attribs['axis']] for input in op.inputs], dtype=np.int32), dtype=DTYPE_INT32), remove=False) return (output_shape, [len(op.inputs)]), (op.inputs[0].dtype, DTYPE_INT32)
def transform_io(g, io_transform, transforms_by_name, driver): # type:(BaseGraph, TrafoOrTrafoDictType, typing.Dict[str, typing.List[Transform]], DataFormatOptimizationDriver)->None io_tensors_by_name = {t.name: t for t in list(g.inputs) + list(g.outputs)} transform_by_io_tensor = {} if isinstance(io_transform, dict): for k, v in six.iteritems(io_transform): assert isinstance(k, (str, driver.tensor_type)), \ "io_transform: Key type must be {} or str".format(driver.tensor_type.__name__) assert isinstance( v, Transform), "io_transform: Value type must be Transform" if isinstance(k, BaseTensor): assert k in six.itervalues(io_tensors_by_name) else: assert k in io_tensors_by_name k = io_tensors_by_name[k] transform_by_io_tensor[k] = v for io_tensor in six.itervalues(io_tensors_by_name): assert io_tensor in transform_by_io_tensor, \ "io_transform: Please specify transform for all io tensors. " \ "You can use graph_optimizer.IDENTITY if no change is required." else: assert isinstance(io_transform, Transform), \ "io_transform must be Transform or Dict[str, Transform] or Dict[NNEFTensor, Transform]" for t in six.itervalues(io_tensors_by_name): transform_by_io_tensor[t] = io_transform for tensor, transform in six.iteritems(transform_by_io_tensor): assert bool(tensor in g.inputs) != bool(tensor in g.outputs), \ "Tensor must be input or output (and not both)" assert isinstance(transform, (Transpose, _CustomTransform)), \ "Unsupported io_transform" if isinstance(transform, _Identity): continue if isinstance(transform, _SmartTFNCHWToNCHW): try: _transform_tf_filter_grad_to_nnef(g, tensor, transforms_by_name, driver) except _TransformException: pass continue if isinstance(transform, _SmartNHWCToNCHW): if tensor.rank <= 2: continue transform = Transpose([0, tensor.rank - 1] + list(range(tensor.rank))[1:-1]) if isinstance(transform, _SmartTFNHWCToNCHW): try: _transform_tf_filter_grad_to_nnef(g, tensor, transforms_by_name, driver) continue except _TransformException: if tensor.rank <= 2: continue transform = Transpose([0, tensor.rank - 1] + list(range(tensor.rank))[1:-1]) if isinstance(transform, _SmartNCHWToTFNCHW): try: _transform_nnef_filter_grad_to_tf(g, tensor, transforms_by_name, driver) except _TransformException: pass continue if isinstance(transform, _SmartNCHWToNHWC): if tensor.rank <= 2: continue transform = Transpose([0] + list(range(tensor.rank))[2:] + [1]) if isinstance(transform, _SmartNCHWToTFNHWC): try: _transform_nnef_filter_grad_to_tf(g, tensor, transforms_by_name, driver) continue except _TransformException: if tensor.rank <= 2: continue transform = Transpose([0] + list(range(tensor.rank))[2:] + [1]) if isinstance(transform, _TFFilterGradToNNEF): _transform_tf_filter_grad_to_nnef(g, tensor, transforms_by_name, driver) continue if isinstance(transform, _NNEFFilterGradToTF): _transform_nnef_filter_grad_to_tf(g, tensor, transforms_by_name, driver) continue assert isinstance( transform, Transpose), "Unsupported io_transform: {}".format(transform) assert len(transform.axes) == tensor.rank, "Transpose: invalid rank" if transform.is_identity(): continue if tensor in g.inputs: assert tensor.name new_input_tensor = driver.create_tensor( graph=g, name=tensor.name, shape=utils.apply_permutation(tensor.shape, transform.axes), dtype=tensor.dtype) add_transform(transforms_by_name, new_input_tensor, transform) transpose = driver.create_transpose_op( graph=g, input=new_input_tensor, axes=utils.inverse_permutation(transform.axes), output=driver.create_tensor(graph=g, name=None, shape=tensor.shape, dtype=tensor.dtype)) graph_utils.replace_tensor_in_inputs(g, tensor, new_input_tensor) graph_utils.replace_tensor_in_consumers(g, tensor, transpose.output, remove=True) else: # output transpose = driver.create_transpose_op( graph=g, input=tensor, axes=transform.axes, output=driver.create_tensor(graph=g, name=tensor.name, shape=utils.apply_permutation( tensor.shape, transform.axes), dtype=tensor.dtype)) add_transform(transforms_by_name, transpose.output, transform) tensor.name = None graph_utils.replace_tensor_in_outputs(g, tensor, transpose.output)
def transform_fused_batch_norm(g, op): # type: (TFGraph, TFOperation)->None VARIANCE_CORRECTION_ENABLED = True in_input = op.inputs[0] in_scale = op.inputs[1] in_offset = op.inputs[2] epsilon = op.attribs["epsilon"] out_y = op.outputs[0] out_batch_mean = op.outputs[1] out_batch_var = op.outputs[2] data_format = op.attribs["data_format"].upper( ) if op.attribs["data_format"] else "NHWC" channel_dim = 1 if data_format == "NCHW" else in_input.rank - 1 rest_count = int(op.inputs[0].count / channel_dim) tensors_to_remove = [] if op.attribs["is_training"]: if VARIANCE_CORRECTION_ENABLED: biased_batch_var = TFTensor(graph=g, shape=list(out_batch_var.shape), dtype=out_batch_var.dtype) const = TFTensor(graph=g, shape=[], dtype=in_input.dtype, data=float(rest_count) / max(rest_count - 1, 1)) TFOperation(graph=g, name="tf.nn.moments", inputs=in_input, attribs=dict(axes=utils.without( range(in_input.rank), channel_dim), keep_dims=False), outputs=(out_batch_mean, biased_batch_var)) TFOperation(graph=g, name="tf.multiply", inputs=(biased_batch_var, const), outputs=out_batch_var) TFOperation(graph=g, name="tf.nn.batch_normalization", inputs=(in_input, out_batch_mean, out_batch_var, in_offset, in_scale), attribs=dict(variance_epsilon=epsilon, _data_format=data_format), outputs=out_y) if len(op.outputs) > 3: # This can happen in gradients out_saved_mean = op.outputs[3] out_saved_var = op.outputs[4] graph_utils.replace_tensor_in_consumers( g, out_saved_mean, out_batch_mean) graph_utils.replace_tensor_in_consumers( g, out_saved_var, out_batch_var) tensors_to_remove += [out_saved_mean, out_saved_var] else: # not VARIANCE_CORRECTION_ENABLED TFOperation(graph=g, name="tf.nn.moments", inputs=in_input, attribs=dict(axes=utils.without( range(in_input.rank), channel_dim), keep_dims=False), outputs=(out_batch_mean, out_batch_var)) TFOperation(graph=g, name="tf.nn.batch_normalization", inputs=(in_input, out_batch_mean, out_batch_var, in_offset, in_scale), attribs=dict(variance_epsilon=epsilon, _data_format=data_format), outputs=out_y) if len(op.outputs) > 3: # This can happen in gradients out_saved_mean = op.outputs[3] out_saved_var = op.outputs[4] graph_utils.replace_tensor_in_consumers( g, out_saved_mean, out_batch_mean) graph_utils.replace_tensor_in_consumers( g, out_saved_var, out_batch_var) tensors_to_remove += [out_saved_mean, out_saved_var] else: # not training in_mean = op.inputs[3] in_variance = op.inputs[4] graph_utils.replace_tensor_in_consumers(g, out_batch_mean, in_mean) graph_utils.replace_tensor_in_consumers(g, out_batch_var, in_variance) tensors_to_remove += [out_batch_mean, out_batch_var] if len(op.outputs) > 3: # This can happen in gradients out_saved_mean = op.outputs[3] out_saved_var = op.outputs[4] graph_utils.replace_tensor_in_consumers(g, out_saved_mean, in_mean) graph_utils.replace_tensor_in_consumers(g, out_saved_var, in_variance) tensors_to_remove += [out_saved_mean, out_saved_var] TFOperation(graph=g, name="tf.nn.batch_normalization", inputs=(in_input, in_mean, in_variance, in_offset, in_scale), attribs=dict(variance_epsilon=epsilon, _data_format=data_format), outputs=out_y) g.remove_operation(op, unlink=True) for t in tensors_to_remove: g.remove_tensor(t)