def testSequencePreprocessor(self): sub1_p = input_preprocessors.ConstantPreprocessor.Params().Set( name='sub1', constants={'foo': 1}) sub2_p = input_preprocessors.ConstantPreprocessor.Params().Set( name='sub2', constants={'bar': 2}) preprocessor_p = input_preprocessors.Sequence.Params().Set( name='list', preprocessors=[sub1_p, sub2_p]) features = py_utils.NestedMap() shapes = py_utils.NestedMap() dtypes = py_utils.NestedMap() preprocessor = preprocessor_p.Instantiate() new_features = preprocessor.TransformFeatures(features) new_shapes = preprocessor.TransformShapes(shapes) new_dtypes = preprocessor.TransformDTypes(dtypes) # Verify shape and dtype self.assertEqual(new_shapes.foo, tf.TensorShape([])) self.assertEqual(new_shapes.bar, tf.TensorShape([])) self.assertEqual(new_dtypes.foo, tf.int64) self.assertEqual(new_dtypes.bar, tf.int64) with self.session() as sess: np_new_features = sess.run(new_features) # Check the new constants exist in the features for both preprocessors self.assertEqual(np_new_features.foo, 1) self.assertEqual(np_new_features.bar, 2)
def testIdentityPreprocessor(self): input_p = input_preprocessors.ConstantPreprocessor.Params().Set( constants={ 'value1': 1, 'value2': np.array([2]) }) identity_p = input_preprocessors.IdentityPreprocessor.Params() features = py_utils.NestedMap() shapes = py_utils.NestedMap() dtypes = py_utils.NestedMap() preprocessors = [input_p.Instantiate(), identity_p.Instantiate()] for preprocessor in preprocessors: # Verify shape / dtypes. shapes = preprocessor.TransformShapes(shapes) dtypes = preprocessor.TransformDTypes(dtypes) features = preprocessor.TransformFeatures(features) self.assertEqual(self.evaluate(features.value1), 1) self.assertEqual(shapes.value1, tf.TensorShape([])) self.assertEqual(dtypes.value1, tf.int64) self.assertEqual(self.evaluate(features.value2), [2]) self.assertEqual(shapes.value2, tf.TensorShape([1])) self.assertEqual(dtypes.value2, tf.int64)
def Shape(self): """Shape of BBoxes.""" p = self.params shapes = {} for laser in p.cbr_laser_names: cbr_shape = p.cbr_ri_shape[:-1] for returns in p.returns: shape_dict = py_utils.NestedMap({ 'xyz': tf.TensorShape(cbr_shape + [3]), 'features': tf.TensorShape(cbr_shape + [4]), 'mask': tf.TensorShape(cbr_shape), }) shapes['%s_%s' % (laser, returns)] = shape_dict shapes['%s_extrinsics' % laser] = tf.TensorShape([4, 4]) shapes['%s_beam_inclinations' % laser] = tf.TensorShape([2]) for laser in p.gbr_laser_names: gbr_shape = p.gbr_ri_shape[:-1] for returns in p.returns: shape_dict = py_utils.NestedMap({ 'xyz': tf.TensorShape(gbr_shape + [3]), 'features': tf.TensorShape(gbr_shape + [4]), 'mask': tf.TensorShape(gbr_shape), }) shapes['%s_%s' % (laser, returns)] = shape_dict shapes['%s_extrinsics' % laser] = tf.TensorShape([4, 4]) shapes['%s_beam_inclinations' % laser] = tf.TensorShape([64]) return py_utils.NestedMap(shapes)
def Shape(self): p = self.params ret = py_utils.NestedMap( points_xyz=tf.TensorShape([p.max_num_points, 3]), points_feature=tf.TensorShape([p.max_num_points, p.num_features])) if p.max_num_points is not None: ret.points_padding = tf.TensorShape([p.max_num_points]) return ret
def Shape(self): """Shape of images.""" p = self.params shapes = {'frame_pose': tf.TensorShape([4, 4])} for camera_name in p.camera_names: shapes['%s' % camera_name] = tf.TensorShape(p.image_shape) # 1d Array of [f_u, f_v, c_u, c_v, k{1, 2}, p{1, 2}, k{3}]. # Note that this intrinsic corresponds to the images after scaling. # Camera model: pinhole camera. # Lens distortion: # Radial distortion coefficients: k1, k2, k3. # Tangential distortion coefficients: p1, p2. # k_{1, 2, 3}, p_{1, 2} follows the same definition as OpenCV. shapes['%s_intrinsics' % camera_name] = tf.TensorShape([9]) shapes['%s_extrinsics' % camera_name] = tf.TensorShape([4, 4]) shapes['%s_pose' % camera_name] = tf.TensorShape([4, 4]) shapes['%s_velocity' % camera_name] = tf.TensorShape([6]) for feat in [ 'pose_timestamp', 'shutter', 'camera_trigger_time', 'camera_readout_done_time' ]: shapes['%s_%s' % (camera_name, feat)] = tf.TensorShape([]) shapes['%s_rolling_shutter_direction' % camera_name] = tf.TensorShape([]) return py_utils.NestedMap(shapes)
def __init__(self, dtype, shape, send_device, recv_device, name=None): """Construct a channel. Args: dtype: The dtype of tensors sent through the channel. shape: The shape of tensors sent through the channel. Must be a fully defined shape for TPUs. send_device: A fully-specified tensorflow device. recv_device: A fully-specified tensorflow device. name: A name for the channel (optional). """ current_graph = tf.get_default_graph() assert current_graph, "A channel is scoped within a tf.Graph" self._dtype = dtype self._send_device = send_device self._recv_device = recv_device self._name = current_graph.unique_name(name if name else "channel") assert shape is not None shape = tf.TensorShape(shape) self._shape = shape self._send_tpu_core = _TpuCore(send_device) self._recv_tpu_core = _TpuCore(recv_device) self._send_called = False self._recv_op = None assert ((self._send_tpu_core == -1) == (self._recv_tpu_core == -1)), ( "Mixing TPU and non-TPU: %s and %s" % (send_device, recv_device)) if self._send_tpu_core >= 0: assert self._shape.is_fully_defined(), ( "TPU channel must have fully defined shape. Name: %s, shape: %s" % (self._name, self._shape)) assert self._send_tpu_core != self._recv_tpu_core, ( "TPU send/recv must be cross-core: %s and %s" % (send_device, recv_device))
def EncodeModality(self, modality: str, inputs, batch_shape=None): """Runs `inputs` through `modality`'s encoder and optional projection. Args: modality: Name of the modality to encode. inputs: Tensor(s) of input to the encoder, e.g. a batch of decoded images. batch_shape: TensorShape describing the batch structure of the inputs. Defaults to `[None]`, which means `inputs` have a single batch dimension. Set to (e.g.) `[None, 5]` if each example in the batch contains 5 encodable items. Returns: A float32 Tensor of the encoded items, shape `batch_shape + [joint_embedding_dim]` """ if batch_shape is None: batch_shape = tf.TensorShape([None]) if not isinstance(inputs, tuple): inputs = (inputs, ) encodings = _EncodeBatch(self.encoders[modality], inputs, batch_shape=batch_shape) # If necessary, project outputs to joint_embedding_dim. if modality in self.projections: return self.projections[modality](encodings) else: return encodings
def Shape(self): """Shape of BBoxes.""" p = self.params shapes = { 'labels': tf.TensorShape([p.max_num_objects]), 'label_ids': tf.TensorShape([p.max_num_objects]), 'detection_difficulties': tf.TensorShape([p.max_num_objects]), 'tracking_difficulties': tf.TensorShape([p.max_num_objects]), 'bboxes_3d': tf.TensorShape([p.max_num_objects, 7]), 'bboxes_3d_mask': tf.TensorShape([p.max_num_objects]), 'bboxes_3d_num_points': tf.TensorShape([p.max_num_objects]), 'unfiltered_bboxes_3d_mask': tf.TensorShape([p.max_num_objects]), 'speed': tf.TensorShape([p.max_num_objects, 2]), 'acceleration': tf.TensorShape([p.max_num_objects, 2]) } return py_utils.NestedMap(shapes)
def FProp(self, theta, *args): """Runs p.repeat copies of self.body.FProp independently. Args: theta: Layer model parameters. The shape of each variable in theta is always [p.repeat, ...]. And the i-th slice theta[i] becomes theta of the i-th copy of self.body. *args: Input arguments. The shape of each tensor in args is always [p.repeat, ....]. And the list [arg[i] for arg in args] becomes inputs to the i-th copy of self.body.FProp. Returns: The accumulated output_tensors. Each tensor t in the return has the shape [p.repeat, ....] and the tuple (t[i] for i in output_tensors) is the return tuple of the i-th self.body.FProp. """ p = self.params for arg in args: if arg is not None: arg = py_utils.HasShape(arg, [p.repeat], ndims=1) theta_stack = _MaybeStackExtraTheta(theta.body, self.body.vars, p.repeat) inputs = py_utils.NestedMap(theta=theta_stack, args=list(args)) # Infer out_shapes from FPropMeta. out_shapes = self._InferOutShapes(args) def _CellFn(unused_theta, unused_state0, inputs): """Recurrent cell function wrapper of body.FProp.""" # Sets shapes for both theta and inputs to self.body.FProp. for dst, src in zip(inputs.args + inputs.theta.Flatten(), list(args) + theta_stack.Flatten()): if src is not None: dst.set_shape(tf.TensorShape(src.shape.as_list()[1:])) # Runs the actual body.FProp fprop_outputs = self.body.FProp(inputs.theta, *inputs.args) fprop_outputs = _ToTuple(fprop_outputs) assert len(fprop_outputs) == len(out_shapes) # Passes fprop outputs to the next layer through state. state1 = py_utils.NestedMap(outputs=list(fprop_outputs)) return state1, py_utils.NestedMap() with tf.name_scope(p.name): # Initiate state0 with inferred output shapes. state0 = py_utils.NestedMap( outputs=[tf.zeros(shape, args[0].dtype) for shape in out_shapes]) # Runs body.FProp p.repeat times using Recurrent. acc_states, _ = recurrent.Recurrent( theta=py_utils.NestedMap(), state0=state0, inputs=inputs, cell_fn=_CellFn) # Retrieves fprop outputs from state1 and sets shapes. output_tensors = tuple(acc_states.outputs) for out_idx in range(len(output_tensors)): output_tensors[out_idx].set_shape( tf.TensorShape([p.repeat] + out_shapes[out_idx].as_list())) return output_tensors[0] if len(args) == 1 else tuple(output_tensors)
def _get_input_shapes(self, *args): p = self.params if p.nested_map_fprop: assert len(args) == 1 assert isinstance(args[0], py_utils.NestedMap) input_tensors = py_utils.Flatten(args[0]) else: input_tensors = _ToTuple(args) # Get batch size from the first tensor which is not None. mini_batch_size = None for input_tensor in input_tensors: if input_tensor is not None: mini_batch_size = input_tensor.get_shape().as_list()[p.batch_dim] assert mini_batch_size is not None micro_batch_size = p.micro_batch_size if not micro_batch_size: if p.num_micro_batches > mini_batch_size: p.num_micro_batches = mini_batch_size micro_batch_size = mini_batch_size // p.num_micro_batches if mini_batch_size is not None: if micro_batch_size * p.num_micro_batches != mini_batch_size: raise ValueError('micro_batch_size * num_micro_batches != batch_size.') input_shapes = () for input_tensor in input_tensors: if input_tensor is not None: input_shape = input_tensor.get_shape().as_list() input_shape[p.batch_dim] = micro_batch_size input_shapes += (tf.TensorShape(input_shape),) else: input_shapes += (None,) if p.nested_map_fprop: input_shapes = py_utils.Pack(args[0], input_shapes) return input_shapes
def testRandomChoicePreprocessor(self): p = input_preprocessors.RandomChoicePreprocessor.Params() p.weight_tensor_key = 'weights' # Construct 4 preprocessors each producing a different value. p.subprocessors = [ input_preprocessors.ConstantPreprocessor.Params().Set( constants={'value': 1}), input_preprocessors.ConstantPreprocessor.Params().Set( constants={'value': 2}), input_preprocessors.ConstantPreprocessor.Params().Set( constants={'value': 3}), input_preprocessors.ConstantPreprocessor.Params().Set( constants={'value': 4}), ] preprocessor = p.Instantiate() # Construct test data. features = py_utils.NestedMap() features.weights = tf.constant([1., 2., 3., 4.]) shapes = py_utils.NestedMap() shapes.weights = tf.TensorShape([4]) dtypes = py_utils.NestedMap() dtypes.weights = tf.float32 # Verify shape / dtypes. new_shapes = preprocessor.TransformShapes(shapes) new_dtypes = preprocessor.TransformDTypes(dtypes) self.assertEqual(new_shapes.value, tf.TensorShape([])) self.assertEqual(new_dtypes.value, tf.int64) new_features = preprocessor.TransformFeatures(features) counts = [0, 0, 0, 0] with self.session() as sess: # Run 10000 times to get probability distribution. for _ in range(10000): new_features_np = sess.run(new_features) counts[new_features_np.value - 1] += 1 # Check distribution roughly matches [0.1, 0.2, 0.3, 0.4] self.assertTrue(counts[0] > 800 and counts[0] < 1200) self.assertTrue(counts[1] > 1800 and counts[1] < 2200) self.assertTrue(counts[2] > 2800 and counts[2] < 3200) self.assertTrue(counts[3] > 3800 and counts[3] < 4200)
def ToTensorShape(self): """Converts to a possibly partially specified tf.TensorShape.""" dims = [] for d in self._shape: if d.is_number and d.is_integer: dims.append(int(d)) else: dims.append(None) return tf.TensorShape(dims)
def testRandomChoicePreprocessorErrors(self): p = input_preprocessors.RandomChoicePreprocessor.Params() p.weight_tensor_key = 'weights' # Subprocessors produce different shapes p.subprocessors = [ input_preprocessors.ConstantPreprocessor.Params().Set( constants={'value': 1}), input_preprocessors.ConstantPreprocessor.Params().Set( constants={'value': [2, 3]}), ] preprocessor = p.Instantiate() # Construct test data. shapes = py_utils.NestedMap() shapes.weights = tf.TensorShape([2]) with self.assertRaises(ValueError): preprocessor.TransformShapes(shapes) # Subprocessors produce different keys p.subprocessors = [ input_preprocessors.ConstantPreprocessor.Params().Set( constants={'value': 1}), input_preprocessors.ConstantPreprocessor.Params().Set( constants={'foo': 2}), ] preprocessor = p.Instantiate() # Construct test data. shapes = py_utils.NestedMap() shapes.weights = tf.TensorShape([2]) with self.assertRaises(ValueError): preprocessor.TransformShapes(shapes) # Subprocessors produce different dtypes p.subprocessors = [ input_preprocessors.ConstantPreprocessor.Params().Set( constants={'value': 1}), input_preprocessors.ConstantPreprocessor.Params().Set( constants={'value': 2.}), ] preprocessor = p.Instantiate() # Construct test data. dtypes = py_utils.NestedMap() dtypes.weights = tf.float32 with self.assertRaises(ValueError): preprocessor.TransformDTypes(dtypes)
def _Feature(shape, dtype=tf.float32): shape = tf.TensorShape(shape) if shape.is_fully_defined(): return tf.io.FixedLenFeature(shape, dtype=dtype) else: if shape[:1].is_fully_defined() or not shape[1:].is_fully_defined(): raise ValueError(f'Unsupported sequence shape {shape}') return tf.io.FixedLenSequenceFeature(shape[1:], dtype=dtype, allow_missing=True)
def _GetShapes(tensors, none_shapes=False): """Util for getting nested structure of shapes from structure of tensors. Args: tensors: Structure of Tensors to get shapes for. none_shapes: Returns None shapes if true. Returns: The same structure as tensors but of corresponding `TensorShape` objects. """ shapes = [] for t in tf.nest.flatten(tensors): shape = t.get_shape() if isinstance(t, tf.Tensor) else None if none_shapes: if shape: shapes.append(tf.TensorShape([None] * len(shape))) else: shapes.append(tf.TensorShape(None)) else: shapes.append(tf.TensorShape(shape)) return type(tensors)(tf.nest.pack_sequence_as(tensors, shapes))
def Shape(self): """The expected shape of each field.""" return py_utils.NestedMap(pose=tf.TensorShape([4, 4]), run_segment=tf.TensorShape([]), run_start_offset=tf.TensorShape([]), time_of_day=tf.TensorShape([]), location=tf.TensorShape([]), weather=tf.TensorShape([]))
def testMultipleResultsPerExample(self): # Simple batch of 3 examples with 2 items per example in the result # modality. batch_size = 3 results_per_example = 2 inputs = label_lib.ExamplePairs.WithinBatch( batch=dict(some_feature=tf.range(batch_size)), query_modality='q', result_modality='r') def example_pair_labeler(_): return tf.constant([ [1, 0, 0], [0, 1, X], [0, X, 1], ], dtype=tf.int64) multi_item_labeler = label_lib.MultiItemExampleWrapper( example_pair_labeler, modality_batch_shapes=dict(q=tf.TensorShape([None]), r=tf.TensorShape( [None, results_per_example]))) labels = multi_item_labeler(inputs) self.assertEqual([batch_size, batch_size, results_per_example], labels.shape.as_list()) # [3, 3, 2] expected_labels = [ # pyformat: disable [[1, 1], [0, 0], [0, 0]], [[0, 0], [1, 1], [X, X]], [[0, 0], [X, X], [1, 1]] # pyformat: enable ] self.assertAllEqual(expected_labels, labels)
def _CellFn(unused_theta, unused_state0, inputs): """Recurrent cell function wrapper of body.FProp.""" # Sets shapes for both theta and inputs to self.body.FProp. for dst, src in zip(inputs.args + inputs.theta.Flatten(), list(args) + theta_stack.Flatten()): if src is not None: dst.set_shape(tf.TensorShape(src.shape.as_list()[1:])) # Runs the actual body.FProp fprop_outputs = self.body.FProp(inputs.theta, *inputs.args) fprop_outputs = _ToTuple(fprop_outputs) assert len(fprop_outputs) == len(out_shapes) # Passes fprop outputs to the next layer through state. state1 = py_utils.NestedMap(outputs=list(fprop_outputs)) return state1, py_utils.NestedMap()
def _AssignVar(self, var_op): size = var_op.get_attr('dtype').size shape = tf.TensorShape(var_op.get_attr('shape')) assert self._var_space_pq, ('No ps devices to use.') allocated, device = heapq.heappop(self._var_space_pq) if shape.num_elements() is None: # For vars whose shape aren't known statically, make a constant # estimate to avoid introducing more complexities. var_bytes = 10 * 1024**2 * size else: var_bytes = shape.num_elements() * size allocated += var_bytes heapq.heappush(self._var_space_pq, (allocated, device)) tf.logging.info('Place variable %s on %s %d(+%d)', var_op.name, device, allocated, var_bytes) return device
def Shape(self): p = self.params shape = py_utils.NestedMap(width=tf.TensorShape([1]), height=tf.TensorShape([1]), velo_to_image_plane=tf.TensorShape([3, 4]), velo_to_camera=tf.TensorShape([4, 4]), camera_to_velo=tf.TensorShape([4, 4])) if p.decode_image: shape.image = tf.TensorShape( [self._KITTI_MAX_HEIGHT, self._KITTI_MAX_WIDTH, 3]) return shape
def _AssignVar(self, var_op): size = var_op.get_attr('dtype').size shape = tf.TensorShape(var_op.get_attr('shape')) assert self._var_space_pq, ('No ps devices to use.') allocated, device = heapq.heappop(self._var_space_pq) if shape.num_elements() is None: assert var_op.name.endswith( 'wb/var'), 'Unexpected name pattern: %s' % var_op.name # CuDNN RNN vars shape aren't known statically, decide to make a constant # estimate to avoid introducing more complexities. allocated += 10 * 1024**2 * size else: allocated += shape.num_elements() * size heapq.heappush(self._var_space_pq, (allocated, device)) tf.logging.info('Place variable %s on %s %d', var_op.name, device, allocated) return device
def _CellFn(unused_theta, state0, theta_i): """Recurrent cell function wrapper of body.FProp.""" # Retrieves fprop arguments from state and sets shapes. fprop_inputs = _StateToArgs(state0) # Sets shapes for theta_i as well. for dst, src in zip(theta_i.Flatten(), theta_stack.Flatten()): if src is not None: dst.set_shape(tf.TensorShape(src.shape.as_list()[1:])) # Runs the actual body.FProp fprop_outputs = self._body.FProp(theta_i, *fprop_inputs) fprop_outputs = _ToTuple(fprop_outputs) assert len(fprop_outputs) == len(fprop_inputs) # Passes fprop outputs to the next layer through state. state1 = _ArgsToState(fprop_outputs) return state1, py_utils.NestedMap()
def testRandomChoicePreprocessor(self): p = input_preprocessors.RandomChoicePreprocessor.Params() # Construct 4 preprocessors each producing a different value. base = input_preprocessors.ConstantPreprocessor.Params() c1 = (base.Copy().Set(constants={'value': 1}), schedule.Constant.Params().Set(value=1)) c2 = (base.Copy().Set(constants={'value': 2}), schedule.Constant.Params().Set(value=2)) c3 = (base.Copy().Set(constants={'value': 3}), schedule.Constant.Params().Set(value=3)) c4 = (base.Copy().Set(constants={'value': 4}), schedule.Constant.Params().Set(value=4)) p.subprocessors = [c1, c2, c3, c4] # Create global step because schedules depend on it. _ = py_utils.GetOrCreateGlobalStepVar() preprocessor = p.Instantiate() features = py_utils.NestedMap() shapes = py_utils.NestedMap() dtypes = py_utils.NestedMap() # Verify shape / dtypes. new_shapes = preprocessor.TransformShapes(shapes) new_dtypes = preprocessor.TransformDTypes(dtypes) self.assertEqual(new_shapes.value, tf.TensorShape([])) self.assertEqual(new_dtypes.value, tf.int64) self.evaluate(tf.global_variables_initializer()) new_features = preprocessor.TransformFeatures(features) counts = [0, 0, 0, 0] with self.session() as sess: # Run 10000 times to get probability distribution. for _ in range(10000): new_features_np = sess.run(new_features) counts[new_features_np.value - 1] += 1 # Check distribution roughly matches [0.1, 0.2, 0.3, 0.4] self.assertTrue(counts[0] > 800 and counts[0] < 1200) self.assertTrue(counts[1] > 1800 and counts[1] < 2200) self.assertTrue(counts[2] > 2800 and counts[2] < 3200) self.assertTrue(counts[3] > 3800 and counts[3] < 4200)
def testPSRandomSize(self): p = cluster_factory.Cluster.Params() p.worker.name = '/job:trainer' p.ps.name = '/job:ps' p.ps.replicas = 10 c = cluster_factory.Cluster(p) g = tf.Graph() vs = [] np.random.seed(301) with g.as_default(): with tf.device(c.GetPlacer()): # Creates 200 variables with different sizes. for i in range(200): if i % 13: size = np.random.randint(10000) elif i % 7: size = np.random.randint(100) else: size = np.random.randint(10) vs.append(tf.get_variable('x%d' % i, shape=(size))) sum_all = tf.add_n([tf.reduce_sum(x) for x in vs]) # Computes the total size of variables placed on each device. total_size = {} # device name -> size for v in vs: size = tf.TensorShape(v.op.get_attr('shape')).num_elements() if v.device in total_size: total_size[v.device] += size else: total_size[v.device] = size for (device, allocated) in zip( sorted(total_size), [91701, 91361, 90346, 88738, 87240, 89265, 91944, 92472, 88051, 95053]): self.assertEqual(total_size[device], allocated) self.assertEqual( sum_all.device, cluster.MakeDeviceString( job_name='/job:trainer', replica_id=0, task_id=0, device_name='CPU', device_id=0))
def testIntraModalLabels(self): # Simulate a batch of 4 examples with 2 items each in the 'text' modality. batch_size = 4 items_per_example = 2 modality = 'text' modality_shape = tf.TensorShape([batch_size, items_per_example]) inputs = label_lib.ExamplePairs.WithinBatch( batch=dict(some_feature=tf.range(batch_size)), query_modality=modality, result_modality=modality) def example_pair_labeler(_): return tf.constant([ [1, 0, 0, X], [0, 1, 0, 0], [0, 0, 1, 0], [X, 0, 0, 1], ]) labeler = label_lib.MultiItemExampleWrapper( example_pair_labeler, modality_batch_shapes={modality: modality_shape}) labels = labeler(inputs) self.assertEqual(modality_shape + modality_shape, labels.shape) # The pairwise labels actually have rank 4 (twice the rank of ids), but we # compare them in matrix form for easier inspection. There are 8 items # total. Each should have a positive label for every other item from the # same example. Self-pairs should be ignored (they are neither positive # nor negative pairs), as well as pairs from duplicated examples. self.assertAllEqual([ [X, 1, 0, 0, 0, 0, X, X], [1, X, 0, 0, 0, 0, X, X], [0, 0, X, 1, 0, 0, 0, 0], [0, 0, 1, X, 0, 0, 0, 0], [0, 0, 0, 0, X, 1, 0, 0], [0, 0, 0, 0, 1, X, 0, 0], [X, X, 0, 0, 0, 0, X, 1], [X, X, 0, 0, 0, 0, 1, X], ], tf.reshape(labels, [8, 8]))
def Shape(self): return py_utils.NestedMap({self.KEY_NAME: tf.TensorShape([])})
def Shape(self): return py_utils.NestedMap({'foo': tf.TensorShape([])})
def BeamSearchDecode(self, theta, encoder_outputs, num_hyps_per_beam_override=0, init_beam_search_state=None, pre_beam_search_step_callback=None, post_beam_search_step_callback=None, max_steps=None): """Performs beam-search based decoding. Args: theta: A NestedMap object containing weights' values of the decoder layer and its children layers. encoder_outputs: A NestedMap containing encoder outputs to be passed to the callbacks. num_hyps_per_beam_override: If set to a value <= 0, this parameter is ignored. If set to a value > 0, then this value will be used to override `p.num_hyps_per_beam`. init_beam_search_state: The `InitBeamSearchState` callback. Please refer to the class header comments for more details. pre_beam_search_step_callback: The `PreBeamSearchStepCallback` callback. Please refer to the class header comments for more details. post_beam_search_step_callback: The `PostBeamSearchStepCallback` callback. Please refer to the class header comments for more details. max_steps: maximum beam search steps. If None, use self.params.target_seq_len. Returns: A `BeamSearchDecodeOutput`. """ p = self.params num_hyps_per_beam = p.num_hyps_per_beam if num_hyps_per_beam_override > 0: num_hyps_per_beam = num_hyps_per_beam_override if max_steps is None: max_steps = p.target_seq_len initial_results, other_states = init_beam_search_state( theta, encoder_outputs, num_hyps_per_beam) num_hyps = tf.shape(initial_results.log_probs)[0] num_beams = num_hyps // num_hyps_per_beam if 'step_ids' in initial_results: # [num_hyps, 1] step_ids = tf.ensure_shape(initial_results.step_ids, [None, 1]) else: step_ids = tf.fill([num_hyps, 1], tf.constant(p.target_sos_id, dtype=tf.int32)) min_score = -1e36 best_scores = (tf.zeros(shape=[num_beams], dtype=p.dtype) + min_score) cumulative_scores = tf.zeros(shape=[num_hyps], dtype=p.dtype) in_scores = tf.zeros([max_steps, num_hyps], dtype=p.dtype) in_hyps = tf.zeros([max_steps, num_hyps], dtype=tf.int32) in_prev_hyps = tf.zeros([max_steps, num_hyps], dtype=tf.int32) in_done_hyps = tf.zeros([max_steps, num_hyps], dtype=tf.string) bs_atten_probs = tf.zeros( [max_steps, num_hyps, tf.shape(initial_results.atten_probs)[1]], dtype=p.dtype) cur_step = tf.constant(0, dtype=tf.int32) all_done = tf.constant(False, dtype=tf.bool) core_bs_states = (best_scores, cumulative_scores, in_scores, in_hyps, in_prev_hyps, in_done_hyps, bs_atten_probs) def LoopContinue(cur_step, all_done, unused_step_ids, unused_core_bs_states, unused_other_states_list): return tf.logical_and(cur_step < max_steps, tf.logical_not(all_done)) def LoopBody(cur_step, unused_all_done, step_ids, core_bs_states, other_states_list): (cur_step, all_done, new_step_ids, new_bs_states, new_other_states) = self._BeamSearchStep( theta, encoder_outputs, cur_step, step_ids, core_bs_states, other_states.Pack(other_states_list), num_hyps_per_beam, pre_beam_search_step_callback, post_beam_search_step_callback) return (cur_step, all_done, new_step_ids, new_bs_states, new_other_states.Flatten()) flat_other_states = other_states.Flatten() _, _, _, final_bs_states, flat_final_other_states = tf.while_loop( LoopContinue, LoopBody, loop_vars=(cur_step, all_done, step_ids, core_bs_states, flat_other_states), parallel_iterations=10, back_prop=False, swap_memory=False, shape_invariants=(tf.TensorShape(cur_step.get_shape()), tf.TensorShape(all_done.get_shape()), tf.TensorShape(step_ids.get_shape()), _GetShapes(core_bs_states), _GetShapes(flat_other_states, none_shapes=True))) # [target_seq_len, num_beams * num_hyps_per_beam]. final_done_hyps = final_bs_states[5] final_other_states = other_states.Pack(flat_final_other_states) # TODO(rpang): avoid inspecting 'encoder_outputs'. source_paddings = encoder_outputs.padding if isinstance(source_paddings, py_utils.NestedMap): source_seq_lengths = tf.cast( tf.round( tf.reduce_sum(1.0 - tf.transpose(source_paddings.Flatten()[0]), 1)), tf.int32) else: source_seq_lengths = tf.cast( tf.round(tf.reduce_sum(1.0 - tf.transpose(source_paddings), 1)), tf.int32) # [num_beams, num_hyps_per_beam]. topk_hyps = ops.top_k_terminated_hyps( final_done_hyps, source_seq_lengths, k=num_hyps_per_beam, num_hyps_per_beam=num_hyps_per_beam, length_normalization=p.length_normalization, coverage_penalty=p.coverage_penalty, target_seq_length_ratio=p.target_seq_length_ratio, eoc_id=p.target_eoc_id, merge_paths=p.merge_paths) # [num_beams * num_hyps_per_beam, ...]. max_seq_length = 0 if isinstance(max_steps, tf.Tensor) else max_steps topk_ids, topk_lens, topk_scores = ops.unpack_hyp( tf.reshape(topk_hyps, [-1]), max_seq_length=max_seq_length) # [num_beams, num_hyps_per_beam]. topk_scores = tf.reshape(topk_scores, tf.shape(topk_hyps)) return BeamSearchDecodeOutput(final_done_hyps, topk_hyps, topk_ids, topk_lens, topk_scores, None, final_other_states)
def Transform(self, dataset): """Batches a dataset containing NestedMaps of tensors.""" p = self.params require_sequential_order = p.require_sequential_order or self.do_eval seqlen_fn = getattr(self._input_generator, p.seqlen_fn) def SetBucketKeys(example): example.bucket_keys = seqlen_fn(example) return example dataset = dataset.map(SetBucketKeys, num_parallel_calls=tf.data.experimental.AUTOTUNE, deterministic=require_sequential_order) dataset = dataset.filter( lambda x: x.bucket_keys <= p.bucket_upper_bound[-1]) dataset_structure = py_utils.NestedMap.FromNestedDict( tf.data.experimental.get_structure(dataset)) input_shape_fn = getattr(self._input_generator, p.input_shape_fn) padded_shapes = dataset_structure.TransformWithKey( lambda k, _: tf.TensorShape(input_shape_fn(k))) input_padding_fn = getattr(self._input_generator, p.input_padding_fn) padding_values = dataset_structure.TransformWithKey(input_padding_fn) dataset_structure.VLog(0, 'dataset_structure:') padded_shapes.VLog(0, 'padded_shapes:') bucket_batch_limit = [ batch_utils.scale_split_to_infeed( b, self._input_generator.params.use_per_host_infeed) for b in p.bucket_batch_limit ] dataset = dataset.apply( tf.data.experimental.bucket_by_sequence_length( lambda x: x.bucket_keys, # Upper-bound for bucket_by_sequence_length is exclusive, so add 1 # TODO(jeffreyzhao): There is a off-by-one bug with the upper bound # boundary check, so add 2 instead. Remove when fixed. [x + 2 for x in p.bucket_upper_bound], bucket_batch_limit + [1], padded_shapes=padded_shapes, padding_values=padding_values, pad_to_bucket_boundary=True, drop_remainder=py_utils.use_tpu())) if py_utils.use_tpu(): # Set static shapes for TPU. if min(bucket_batch_limit) != max(bucket_batch_limit): raise ValueError('TPU requires constant batch sizes.') else: b = bucket_batch_limit[0] def SetShape(element): for t in element.Flatten(): t.set_shape((b, ) + t.shape[1:]) return element dataset = dataset.map( SetShape, num_parallel_calls=tf.data.experimental.AUTOTUNE, deterministic=require_sequential_order) return dataset
def __init__(self, batch_shape): batch_shape = tf.TensorShape(batch_shape) assert batch_shape.rank >= 1 batch_shape[1:].assert_is_fully_defined() self._batch_shape = batch_shape self._is_no_op = batch_shape == tf.TensorShape([None])