def _testPad(self, inputs, block_shape, paddings, outputs): block_shape = np.array(block_shape) paddings = np.array(paddings).reshape((len(block_shape), 2)) with self.test_session() as sess, self.test_scope(): for dtype in self.float_types: # TODO(b/68813416): Skip bfloat16's as the input type for direct is # float32 and results in a mismatch, while making testDirect provide the # correctly typed input results in 'no fill-function for data-type' # error. if dtype == dtypes.bfloat16.as_numpy_dtype: continue if dtype == np.float16: actual_inputs = np.array(inputs).astype(dtype) actual_paddings = np.array(paddings).astype(dtype) expected_outputs = np.array(outputs).astype(dtype) else: actual_inputs = inputs actual_paddings = paddings expected_outputs = outputs placeholder = array_ops.placeholder(dtype) # outputs = space_to_batch(inputs) x_tf = array_ops.space_to_batch_nd(placeholder, block_shape, actual_paddings) self.assertAllEqual( sess.run(x_tf, {placeholder: actual_inputs}), expected_outputs) # inputs = batch_to_space(outputs) placeholder = array_ops.placeholder(dtype) x_tf = array_ops.batch_to_space_nd(placeholder, block_shape, actual_paddings) self.assertAllEqual( sess.run(x_tf, {placeholder: expected_outputs}), actual_inputs)
def _testStaticShape(self, input_shape, block_shape, paddings, error): block_shape = np.array(block_shape) paddings = np.array(paddings) # Try with sizes known at graph construction time. with self.assertRaises(error): _ = array_ops.batch_to_space_nd( np.zeros(input_shape, np.float32), block_shape, paddings)
def _testPad(self, inputs, block_shape, paddings, outputs): block_shape = np.array(block_shape) paddings = np.array(paddings).reshape((len(block_shape), 2)) for use_gpu in [False, True]: with self.test_session(use_gpu=use_gpu): # outputs = space_to_batch(inputs) x_tf = array_ops.space_to_batch_nd( math_ops.to_float(inputs), block_shape, paddings) self.assertAllEqual(x_tf.eval(), outputs) # inputs = batch_to_space(outputs) x_tf = array_ops.batch_to_space_nd( math_ops.to_float(outputs), block_shape, paddings) self.assertAllEqual(x_tf.eval(), inputs)
def _testPad(self, inputs, block_shape, paddings, outputs): block_shape = np.array(block_shape) paddings = np.array(paddings).reshape((len(block_shape), 2)) with self.test_session() as sess, self.test_scope(): for dtype in self.float_types: placeholder = array_ops.placeholder(dtype) # outputs = space_to_batch(inputs) x_tf = array_ops.space_to_batch_nd(placeholder, block_shape, paddings) self.assertAllEqual(sess.run(x_tf, {placeholder: inputs}), outputs) # inputs = batch_to_space(outputs) placeholder = array_ops.placeholder(dtype) x_tf = array_ops.batch_to_space_nd(placeholder, block_shape, paddings) self.assertAllEqual(sess.run(x_tf, {placeholder: outputs}), inputs)
def _checkGrad(self, x, block_shape, crops): block_shape = np.array(block_shape) crops = np.array(crops).reshape((len(block_shape), 2)) with self.test_session(): tf_x = ops.convert_to_tensor(x) tf_y = array_ops.batch_to_space_nd(tf_x, block_shape, crops) epsilon = 1e-5 ((x_jacob_t, x_jacob_n)) = gradient_checker.compute_gradient( tf_x, x.shape, tf_y, tf_y.get_shape().as_list(), x_init_value=x, delta=epsilon) self.assertAllClose(x_jacob_t, x_jacob_n, rtol=1e-2, atol=epsilon)
def _testDynamicShape(self, input_shape, block_shape, paddings): block_shape = np.array(block_shape) paddings = np.array(paddings) # Try with sizes unknown at graph construction time. input_placeholder = array_ops.placeholder(dtypes.float32) block_shape_placeholder = array_ops.placeholder( dtypes.int32, shape=block_shape.shape) paddings_placeholder = array_ops.placeholder(dtypes.int32) t = array_ops.batch_to_space_nd(input_placeholder, block_shape_placeholder, paddings_placeholder) with self.assertRaises(ValueError): _ = t.eval({ input_placeholder: np.zeros(input_shape, np.float32), block_shape_placeholder: block_shape, paddings_placeholder: paddings })
def _CloneWithNewOperands(layer_op, input_tensor, weight_tensor, batch_to_space_op): """Clones layer_op with input_tensor and weight_tensor as new inputs.""" new_layer_name = layer_op.name.split('/')[-1] + '_Fold' if layer_op.type == 'Conv2D': return nn_ops.conv2d( input_tensor, weight_tensor, strides=layer_op.get_attr('strides'), padding=layer_op.get_attr('padding'), use_cudnn_on_gpu=layer_op.get_attr('use_cudnn_on_gpu'), data_format=layer_op.get_attr('data_format'), name=new_layer_name) elif layer_op.type == 'MatMul': return math_ops.matmul( input_tensor, weight_tensor, transpose_a=layer_op.get_attr('transpose_a'), transpose_b=layer_op.get_attr('transpose_b'), name=new_layer_name) elif layer_op.type == 'DepthwiseConv2dNative': # We don't copy dilation rate because we reuse the input SpaceToBatch # and create our own BatchToSpace operation below. conv = nn.depthwise_conv2d( input_tensor, weight_tensor, strides=layer_op.get_attr('strides'), padding=layer_op.get_attr('padding'), name=new_layer_name) # Copy the batch to space operation if we have a atrous convolution. if batch_to_space_op: batch_to_space_op = layer_op.outputs[0].consumers()[0] # TODO(suharshs): It's hard to make this name match with the unfused name. # Restructure this code to not rely on scope at all. new_batch_to_space_name = batch_to_space_op.name.split('/')[-1] + '_Fold' conv = array_ops.batch_to_space_nd( conv, batch_to_space_op.inputs[1], batch_to_space_op.inputs[2], name=new_batch_to_space_name) return conv else: raise ValueError('Cannot handle operation of type: %s' % layer_op.type)
def testUnknownShape(self): # Verify that input shape and paddings shape can be unknown. _ = array_ops.batch_to_space_nd( array_ops.placeholder(dtypes.float32), array_ops.placeholder( dtypes.int32, shape=(2,)), array_ops.placeholder(dtypes.int32)) # Only number of input dimensions is known. t = array_ops.batch_to_space_nd( array_ops.placeholder( dtypes.float32, shape=(None, None, None, None)), array_ops.placeholder( dtypes.int32, shape=(2,)), array_ops.placeholder(dtypes.int32)) self.assertEqual(4, t.get_shape().ndims) # Dimensions are partially known. t = array_ops.batch_to_space_nd( array_ops.placeholder( dtypes.float32, shape=(None, None, None, 2)), array_ops.placeholder( dtypes.int32, shape=(2,)), array_ops.placeholder(dtypes.int32)) self.assertEqual([None, None, None, 2], t.get_shape().as_list()) # Dimensions are partially known. t = array_ops.batch_to_space_nd( array_ops.placeholder( dtypes.float32, shape=(3 * 2 * 3, None, None, 2)), [2, 3], array_ops.placeholder(dtypes.int32)) self.assertEqual([3, None, None, 2], t.get_shape().as_list()) # Dimensions are partially known. t = array_ops.batch_to_space_nd( array_ops.placeholder( dtypes.float32, shape=(3 * 2 * 3, None, 2, 2)), [2, 3], [[1, 1], [0, 1]]) self.assertEqual([3, None, 5, 2], t.get_shape().as_list()) # Dimensions are fully known. t = array_ops.batch_to_space_nd( array_ops.placeholder( dtypes.float32, shape=(3 * 2 * 3, 2, 1, 2)), [2, 3], [[1, 1], [0, 0]]) self.assertEqual([3, 2, 3, 2], t.get_shape().as_list())
def _SpaceToBatchNDGrad(op, grad): # Its gradient is the opposite op: BatchToSpaceND. return [array_ops.batch_to_space_nd(grad, op.inputs[1], op.inputs[2]), None, None]
def _CreateFoldedOp(graph, context, has_scaling, freeze_batch_norm_delay, is_training): """Folds in batch norm layer into preceding convolution or FC layer. Creates 3 new nodes, connects their inputs and adds them to the graph: mul is cloned into mul_fold, Conv2D or MatMul, or DepthwiseConv2d is cloned into respective *_Fold, add is cloned into add_fold. Args: graph: Graph to modify. context: String, batch norm context, i.e. node into which BatchNorm is nested. has_scaling: Whether the batch norm has scaling enabled. freeze_batch_norm_delay: How many steps to wait before freezing moving mean and variance and using them for batch normalization. is_training: Bool, true if training. Raises: ValueError: When operation type is not supported, or input and output tensor shapes mismatch for created operations: mul_fold, add_fold. Returns: A pair of Operations, the first is the original consumer node of the batch norm (../BatchNorm/batchnorm_1/add_1), the second is the consumer node of the folded graph (add_fold). """ mul_scale_name = 'mul_1' if has_scaling else 'mul' mul_scale = graph.get_operation_by_name(context + 'BatchNorm/batchnorm_1/' + mul_scale_name) op_below = mul_scale.inputs[0].op # Skip over the BatchToSpace operation in the case of atrous convolutions. batch_to_space_op = None if op_below.type == 'BatchToSpaceND': batch_to_space_op = op_below op_below = op_below.inputs[0].op weights = op_below.inputs[1] match = _GetBatchNormParams( graph=graph, context=context, has_scaling=has_scaling) correction_scale, correction_recip, correction_offset = None, None, None if is_training: correction_scale, correction_recip, correction_offset = ( _ComputeBatchNormCorrections( context=context, match=match, freeze_batch_norm_delay=freeze_batch_norm_delay)) # Special handling for weights of depthwise convolution. if op_below.type == 'DepthwiseConv2dNative': new_shape = [ weights.get_shape().as_list()[2], weights.get_shape().as_list()[3] ] scale_name = 'mul' if has_scaling else 'Rsqrt' scale = graph.get_operation_by_name(context + 'BatchNorm/batchnorm_1/' + scale_name) scale = array_ops.reshape(scale.outputs[0], new_shape, context + 'scale_reshape') if correction_scale is not None: correction_scale = array_ops.reshape(correction_scale, new_shape, context + 'correction_reshape') with ops.device(mul_scale.device): weights = math_ops.multiply(correction_scale, weights, context + 'correction_mult') mul_fold = _CloneOp(mul_scale, context + 'mul_fold', [(0, weights), (1, scale)]) elif op_below.type in ['Conv2D', 'MatMul']: if correction_scale is not None: with ops.device(mul_scale.device): weights = math_ops.multiply(correction_scale, weights, context + 'correction_mult') mul_fold = _CloneOp(mul_scale, context + 'mul_fold', [(0, weights)]) else: raise ValueError('Cannot handle operation of type: %s' % op_below.type) _AssertShapesMatch('mul_fold', mul_fold.inputs[0], mul_fold.outputs[0]) conv_or_fc_folded = _CloneOp(op_below, op_below.name + '_Fold', [(1, mul_fold.outputs[0])]) add_shift = graph.get_operation_by_name(context + 'BatchNorm/batchnorm_1/add_1') corrected_output = conv_or_fc_folded.outputs[0] # Copy the batch to space operation if we have a atrous convolution. if batch_to_space_op: corrected_output = array_ops.batch_to_space_nd( corrected_output, batch_to_space_op.inputs[1], batch_to_space_op.inputs[2], name=batch_to_space_op.name + '_Fold') if correction_offset is not None: with ops.device(conv_or_fc_folded.device): corrected_output = math_ops.multiply(correction_recip, corrected_output, context + 'post_conv_mul') corrected_output = math_ops.add(corrected_output, (correction_offset), context + 'correction_add') add_fold = _CloneOp(add_shift, context + 'add_fold', [(0, corrected_output)]) _AssertShapesMatch('add_fold', add_fold.inputs[0], add_fold.outputs[0]) return add_shift, add_fold
def loop_fn(i): x1 = array_ops.gather(x, i) return array_ops.batch_to_space_nd(x1, block_shapes, crops)
def _SpaceToBatchNDGrad(op, grad): # Its gradient is the opposite op: BatchToSpaceND. return [ array_ops.batch_to_space_nd(grad, op.inputs[1], op.inputs[2]), None, None ]
def _CreateFoldedOp(graph, context, has_scaling, freeze_batch_norm_delay, is_training): """Folds in batch norm layer into preceding convolution or FC layer. Creates 3 new nodes, connects their inputs and adds them to the graph: mul is cloned into mul_fold, Conv2D or MatMul, or DepthwiseConv2d is cloned into respective *_Fold, add is cloned into add_fold. Args: graph: Graph to modify. context: String, batch norm context, i.e. node into which BatchNorm is nested. has_scaling: Whether the batch norm has scaling enabled. freeze_batch_norm_delay: How many steps to wait before freezing moving mean and variance and using them for batch normalization. is_training: Bool, true if training. Raises: ValueError: When operation type is not supported, or input and output tensor shapes mismatch for created operations: mul_fold, add_fold. Returns: A pair of Operations, the first is the original consumer node of the batch norm (../BatchNorm/batchnorm_1/add_1), the second is the consumer node of the folded graph (add_fold). """ mul_scale_name = 'mul_1' if has_scaling else 'mul' mul_scale = graph.get_operation_by_name(context + 'BatchNorm/batchnorm_1/' + mul_scale_name) op_below = mul_scale.inputs[0].op # Skip over the BatchToSpace operation in the case of atrous convolutions. batch_to_space_op = None if op_below.type == 'BatchToSpaceND': batch_to_space_op = op_below op_below = op_below.inputs[0].op weights = op_below.inputs[1] match = _GetBatchNormParams(graph=graph, context=context, has_scaling=has_scaling) correction_scale, correction_recip, correction_offset = None, None, None if is_training: correction_scale, correction_recip, correction_offset = ( _ComputeBatchNormCorrections( context=context, match=match, freeze_batch_norm_delay=freeze_batch_norm_delay)) # Special handling for weights of depthwise convolution. if op_below.type == 'DepthwiseConv2dNative': new_shape = [ weights.get_shape().as_list()[2], weights.get_shape().as_list()[3] ] scale_name = 'mul' if has_scaling else 'Rsqrt' scale = graph.get_operation_by_name(context + 'BatchNorm/batchnorm_1/' + scale_name) scale = array_ops.reshape(scale.outputs[0], new_shape, context + 'scale_reshape') if correction_scale is not None: correction_scale = array_ops.reshape( correction_scale, new_shape, context + 'correction_reshape') with ops.device(mul_scale.device): weights = math_ops.multiply(correction_scale, weights, context + 'correction_mult') mul_fold = _CloneOp(mul_scale, context + 'mul_fold', [(0, weights), (1, scale)]) elif op_below.type in ['Conv2D', 'MatMul']: if correction_scale is not None: with ops.device(mul_scale.device): weights = math_ops.multiply(correction_scale, weights, context + 'correction_mult') mul_fold = _CloneOp(mul_scale, context + 'mul_fold', [(0, weights)]) else: raise ValueError('Cannot handle operation of type: %s' % op_below.type) _AssertShapesMatch('mul_fold', mul_fold.inputs[0], mul_fold.outputs[0]) conv_or_fc_folded = _CloneOp(op_below, op_below.name + '_Fold', [(1, mul_fold.outputs[0])]) add_shift = graph.get_operation_by_name(context + 'BatchNorm/batchnorm_1/add_1') corrected_output = conv_or_fc_folded.outputs[0] # Copy the batch to space operation if we have a atrous convolution. if batch_to_space_op: corrected_output = array_ops.batch_to_space_nd( corrected_output, batch_to_space_op.inputs[1], batch_to_space_op.inputs[2], name=batch_to_space_op.name + '_Fold') if correction_offset is not None: with ops.device(conv_or_fc_folded.device): corrected_output = math_ops.multiply(correction_recip, corrected_output, context + 'post_conv_mul') corrected_output = math_ops.add(corrected_output, (correction_offset), context + 'correction_add') add_fold = _CloneOp(add_shift, context + 'add_fold', [(0, corrected_output)]) _AssertShapesMatch('add_fold', add_fold.inputs[0], add_fold.outputs[0]) return add_shift, add_fold