def FProp(self, theta, inputs, paddings): """Apply convolution to inputs. Args: theta: A `.NestedMap` object containing weights' values of this layer and its children layers. inputs: The inputs tensor. It is expected to be of shape [batch, time, frequency, channel]. The time dimension corresponds to the height dimension as in images and the frequency dimension corresponds to the width dimension as in images. paddings: The paddings tensor, expected to be of shape [batch, time]. Returns: outputs, out_paddings pair. """ p = self.params with tf.name_scope(p.name): inputs = py_utils.with_dependencies([ py_utils.assert_shape_match(tf.shape(paddings), [-1, -1]), py_utils.assert_shape_match( tf.shape(inputs), tf.concat([ tf.shape(paddings), [-1, symbolic.ToStatic(self.input_channels)] ], 0)) ], inputs) def _ApplyPadding(tensor_in, padding_in): padding_expanded = tf.expand_dims(tf.expand_dims(padding_in, -1), -1) return tensor_in * (1.0 - padding_expanded) # Zeroing out padded inputs. inputs = _ApplyPadding(inputs, paddings) # Apply conv on 'inputs'. out = self._ApplyConv(theta, inputs) # NOTE: this may be slightly inaccurate when p.dilation_rate[0] > 1. # But there's likely no real problems. Trying to set it gives an error: # pooling with SAME padding is not implemented for dilation_rate > 1. # NOTE: we use window=p.filter_stride[0] to be compatible with legacy # implementation. Consider updating it to be the actual shape. conv_padding = ComputeConvOutputPadding( paddings, window=p.filter_stride[0], stride=p.filter_stride[0]) # Assuming padded nodes will be properly zero-ed out if necessary by # sub-sequent layers. # out = _ApplyPadding(out, conv_padding) out = py_utils.HasShape( out, symbolic.ToStatic(self.OutShape(tf.shape(inputs)))) return out, conv_padding
def testEvalExpr(self): x = symbolic.Symbol('x') y = symbolic.Symbol('y') xy = x * y # Without symbol-to-value map. self.assertEqual(xy, symbolic.ToStatic(xy)) self.assertEqual(xy, symbolic.ToTensor(xy)) with symbolic.SymbolToValueMap(symbolic.STATIC_VALUES, {x: 2, y: 3}): self.assertEqual(symbolic.ToStatic(xy), 6) # The inner map overrides the outer map. with symbolic.SymbolToValueMap(symbolic.STATIC_VALUES, {x: 5, y: 6}): self.assertEqual(symbolic.ToStatic(xy), 30) # Back to the outer map. self.assertEqual(symbolic.ToStatic(xy), 6) # EvalExpr can also evaluate a symbolic expression to a # Tensor. a = tf.placeholder(tf.float32) b = tf.placeholder(tf.float32) with symbolic.SymbolToValueMap(symbolic.TENSOR_VALUES, {x: a, y: b}): with symbolic.SymbolToValueMap(symbolic.STATIC_VALUES, {x: 2, y: 3}): # Value maps of different types do not affect each other. self.assertEqual(symbolic.ToStatic(xy), 6) ab = symbolic.ToTensor(xy) self.assertIsInstance(ab, tf.Tensor) with self.session() as sess: self.assertEqual(12, sess.run(ab, {a: 3, b: 4})) # EvalExpr supports partial evaluation. with symbolic.SymbolToValueMap(symbolic.STATIC_VALUES, {y: 3}): x3 = symbolic.ToStatic(xy) with symbolic.SymbolToValueMap(symbolic.STATIC_VALUES, {x: 9}): self.assertEqual(27, symbolic.ToStatic(x3))
def testToFromProto(self): outer = hyperparams.Params() outer.Define('integer_val', 1, '') outer.Define('cls_type', type(int), '') inner = hyperparams.Params() inner.Define('float_val', 2.71, '') inner.Define('string_val', 'rosalie et adrien', '') inner.Define('bool_val', True, '') inner.Define('list_of_tuples_of_dicts', [({'string_key': 1729})], '') inner.Define('range', range(1, 3), '') outer.Define('inner', inner, '') outer.Define('empty_list', [], '') outer.Define('empty_tuple', (), '') outer.Define('empty_dict', {}, '') outer.Define('enum', TestEnum.B, '') outer.Define('proto', hyperparams_pb2.HyperparamValue(int_val=42), '') outer.Define('dataclass', TestDataClass(a=[42], b=tf.float32), '') outer.Define('namedtuple', tf.io.FixedLenSequenceFeature([42], tf.float32), '') outer.Define('symbol_x', symbolic.Symbol('x'), '') outer.Define('symbol_2x', outer.symbol_x * 2, '') rebuilt_outer = hyperparams.InstantiableParams.FromProto( outer.ToProto()) self.assertNotIn('cls', rebuilt_outer) self.assertEqual(outer.integer_val, rebuilt_outer.integer_val) self.assertEqual(outer.cls_type, rebuilt_outer.cls_type) self.assertNear(outer.inner.float_val, rebuilt_outer.inner.float_val, 1e-6) self.assertEqual(outer.inner.string_val, rebuilt_outer.inner.string_val) self.assertEqual(outer.inner.bool_val, rebuilt_outer.inner.bool_val) self.assertEqual(outer.inner.list_of_tuples_of_dicts, rebuilt_outer.inner.list_of_tuples_of_dicts) self.assertEqual([1, 2], rebuilt_outer.inner.range) # Rebuilt as list. self.assertEqual(outer.empty_list, rebuilt_outer.empty_list) self.assertEqual(outer.empty_tuple, rebuilt_outer.empty_tuple) self.assertEqual(outer.empty_dict, rebuilt_outer.empty_dict) self.assertEqual(outer.enum, rebuilt_outer.enum) self.assertEqual(outer.proto, rebuilt_outer.proto) self.assertEqual(outer.dataclass, rebuilt_outer.dataclass) self.assertEqual(outer.namedtuple, rebuilt_outer.namedtuple) with symbolic.SymbolToValueMap(symbolic.STATIC_VALUES, {rebuilt_outer.symbol_x: 42}): self.assertEqual(symbolic.ToStatic(rebuilt_outer.symbol_2x), 84)
def GetProjectLastDim(cls, inputs, weight, input_dim, output_dim, proj_obj): """Linear projection on the last dim of the input tensor along with pruning. This is a TPU efficient implementation to avoid reshaping inputs to Rank-2 tensor by using Einsum for the compute. Args: inputs: An input Tensor, the last dimension of which is input_dim. weight: A weight matrix with shape [input_dim, output_dim]. input_dim: An integer or a symbolic dim, the last dimension of the inputs. output_dim: An integer or a symbolic dim, the last dimension of the outputs. proj_obj: a ProjectionLayer object. Returns: An output Tensor of the same rank as inputs, the last dimension is output_dim. """ theta = proj_obj.theta p = proj_obj.params input_dim = int( symbolic.ToStatic(input_dim) if symbolic.IsExpr(input_dim ) else input_dim) output_dim = int( symbolic.ToStatic(output_dim) if symbolic.IsExpr(output_dim ) else output_dim) if (py_utils.use_tpu() and inputs.shape is not None and inputs.shape.rank is not None and inputs.shape.rank < 26): # Avoids reshape if feasible and uses Einsum. if inputs.shape.rank == 2: outputs = tf.matmul(inputs, weight) else: outputs = cls.GetEinSumResult(inputs, proj_obj) else: if p.pruning_hparams_dict[ 'compression_option'] == 9 and p.pruning_hparams_dict[ 'compress_input']: blocked_inputs = tf.reshape( inputs, py_utils.ToStaticShape( [-1, p.pruning_hparams_dict['input_block_size']])) compressed_inputs = tf.reshape( py_utils.Matmul(blocked_inputs, theta.b_matrix_tfvar), py_utils.ToStaticShape([ -1, input_dim // p.pruning_hparams_dict['input_compression_factor'] ])) else: compressed_inputs = tf.reshape( inputs, py_utils.ToStaticShape([-1, input_dim])) if p.pruning_hparams_dict['compression_option'] == 10: if p.pruning_hparams_dict['block_method'] == 'mask': intermediate_result = py_utils.Matmul( compressed_inputs, tf.multiply(theta.c_matrix_tfvar, theta.c_mask_tfvar)) elif p.pruning_hparams_dict['block_method'] == 'loop': num_blocks = p.pruning_hparams_dict[ 'block_compression_factor'] input_splitted = tf.split(compressed_inputs, num_blocks, axis=-1) output_splitted = [] for i, input_i in enumerate(input_splitted): output_splitted.append( py_utils.Matmul(input_i, theta.c_matrix_tfvar[i, :, :])) intermediate_result = tf.concat(output_splitted, axis=-1) else: intermediate_result = py_utils.Matmul(compressed_inputs, theta.c_matrix_tfvar) if p.pruning_hparams_dict[ 'compression_option'] == 9 and p.pruning_hparams_dict[ 'compress_output']: blocked_intermediate_result = tf.reshape( intermediate_result, py_utils.ToStaticShape([ -1, p.pruning_hparams_dict['output_block_size'] // p.pruning_hparams_dict['output_compression_factor'] ])) outputs = py_utils.Matmul(blocked_intermediate_result, theta.d_matrix_tfvar) else: outputs = intermediate_result outputs = tf.reshape( outputs, tf.concat([ tf.cast(py_utils.GetShape(inputs)[:-1], tf.int32), py_utils.ToStaticShape([output_dim]) ], axis=0)) return outputs
def GetProjectLastDim(cls, inputs, weight, input_dim, output_dim, proj_obj): """Linear projection on the last dim of the input tensor along with pruning. This is a TPU efficient implementation to avoid reshaping inputs to Rank-2 tensor by using Einsum for the compute. Args: inputs: An input Tensor, the last dimension of which is input_dim. weight: A weight matrix with shape [input_dim, output_dim]. input_dim: An integer or a symbolic dim, the last dimension of the inputs. output_dim: An integer or a symbolic dim, the last dimension of the outputs. proj_obj: a ProjectionLayer object. Returns: An output Tensor of the same rank as inputs, the last dimension is output_dim. """ theta = proj_obj.theta p = proj_obj.params input_dim = int( symbolic.ToStatic(input_dim) if symbolic.IsExpr(input_dim ) else input_dim) output_dim = int( symbolic.ToStatic(output_dim) if symbolic.IsExpr(output_dim ) else output_dim) if (py_utils.use_tpu() and inputs.shape is not None and inputs.shape.rank is not None and inputs.shape.rank < 26): # Avoids reshape if feasible and uses Einsum. if inputs.shape.rank == 2: outputs = tf.matmul(inputs, weight) else: s = ''.join([chr(x) for x in range(97, 123)]) # abc...xyz r = inputs.shape.rank outputs = cls.GetEinSumResult( inputs, weight, '{0}y,yz->{0}z'.format(s[:r - 1]), proj_obj) else: if p.pruning_hparams_dict['compress_input']: blocked_inputs = tf.reshape( inputs, py_utils.ToStaticShape( [-1, p.pruning_hparams_dict['input_block_size']])) compressed_inputs = tf.reshape( py_utils.Matmul(blocked_inputs, theta.b_matrix_tfvar), py_utils.ToStaticShape([ -1, input_dim // p.pruning_hparams_dict['input_compression_factor'] ])) else: compressed_inputs = tf.reshape( inputs, py_utils.ToStaticShape([-1, input_dim])) intermediate_result = py_utils.Matmul(compressed_inputs, theta.c_matrix_tfvar) if p.pruning_hparams_dict['compress_output']: blocked_intermediate_result = tf.reshape( intermediate_result, py_utils.ToStaticShape([ -1, p.pruning_hparams_dict['output_block_size'] // p.pruning_hparams_dict['output_compression_factor'] ])) outputs = py_utils.Matmul(blocked_intermediate_result, theta.d_matrix_tfvar) else: outputs = intermediate_result outputs = tf.reshape( outputs, tf.concat([ tf.cast(py_utils.GetShape(inputs)[:-1], tf.int32), py_utils.ToStaticShape([output_dim]) ], axis=0)) return outputs