def _TestStreamStepHelper(self, **kwargs): """Main helper method.""" batch_size, max_seqlen, input_dim = 2, 32, kwargs['input_dim'] stride = kwargs.get('stride', 1) # max_seqlen is divisible by stride. assert max_seqlen % stride == 0 right_context = kwargs.get('right_context', 0) # Prepares inputs. inputs, paddings = self._GetInputs(batch_size, max_seqlen, input_dim) # Gets params p = self._GetParams(**kwargs) # Builds graph. with self.session(use_gpu=False) as sess: l = p.Instantiate() init_op = tf.global_variables_initializer() fprop_out = self._FProp(l, inputs, paddings) base_outputs = self._GetFPropOutput(fprop_out) out_rank = py_utils.GetRank(base_outputs) base_outputs *= py_utils.AppendDims(1. - paddings, out_rank - 2) try: state = l.zero_state(batch_size) except TypeError: state = l.zero_state(l.theta, batch_size) outputs = [] for i in range(max_seqlen // stride + int(math.ceil(right_context / stride))): if i < max_seqlen // stride: step_inputs = inputs[:, stride * i:stride * (i + 1)] step_paddings = paddings[:, stride * i:stride * (i + 1)] else: step_inputs = tf.zeros_like(inputs[:, 0:stride]) step_paddings = tf.ones_like(paddings[:, 0:stride]) output, _, state = l.StreamStep(l.theta, step_inputs, step_paddings, state) outputs.append(output) outputs = tf.concat(outputs, axis=1) outputs = self._NormalizeStreamStepOutput(outputs, paddings, right_context, max_seqlen) sess.run(init_op) expected, actual = sess.run([base_outputs, outputs]) print(f'expected: {repr(expected)}, {expected.shape}') print(f'actual: {repr(actual)}, {actual.shape}') print(f'np.sum(np.abs(expected)): {np.sum(np.abs(expected))}') print(f'np.sum(np.abs(actual)): {np.sum(np.abs(actual))}') tol = kwargs.get('tol', 1e-6) self.assertAllClose(expected, actual, atol=tol, rtol=tol)
def _Extract(self, features): """Returns the laser Tensor.""" p = self.params ret = super()._Extract(features) all_vxyz = [] all_classes = [] for lidar in p.lidar_names: for ri in p.lidar_returns: feature_name = 'laser_%s_%s' % (lidar, ri) laser_data = tf.reshape( _Dense(features[feature_name]), [-1, 3 + p.num_features]) num = py_utils.GetShape(laser_data)[0] # We expect lidar_$lidar_$ri and lidar_$lidar_$ri_flow has # same number of points. feature_name += '_flow' laser_data = tf.reshape(_Dense(features[feature_name]), [num, 3 + 1]) points_vxyz = laser_data[..., 0:3] points_classes = laser_data[..., 3] all_vxyz += [points_vxyz] all_classes += [points_classes] # Stack all of the points along the major dimension points_vxyz = tf.concat(all_vxyz, axis=0) points_class = tf.concat(all_classes, axis=0) # The precomputed class uses -1 to mean 5 in our current code. points_class = tf.where( tf.less(points_class, 0), 5. * tf.ones_like(points_class), points_class) if p.max_num_points is not None: assert 'points_padding' in ret points_vxyz = py_utils.PadOrTrimTo(points_vxyz, [p.max_num_points, 3]) points_class = py_utils.PadOrTrimTo(points_class, [p.max_num_points]) assert 'points_xyz' in ret ret.world_flow = points_vxyz ret.pointwise_class = tf.cast(points_class, tf.int32) return ret
def _ProcessMASSInput(self, source_id, src): """Perform MASS input processing.""" if self.do_eval or self.mass_layer is None: # At eval time, we copy src to tgt return self._ProcessSingleInput(source_id, src, src) _, labels, paddings = self.StringsToIds(tf.reshape(src, [1]), is_source=True, key=self._src_tokenizer_key) weights = 1 - paddings actual_seq_len = tf.cast(tf.reduce_sum(weights, 1), tf.int32) src_lang_ids, tgt_lang_ids = self._GetTaskIds(source_id) mass_out = self.mass_layer.Mask(labels, weights, actual_seq_len) features = py_utils.NestedMap() features.src = py_utils.NestedMap() features.src.ids = mass_out.src.ids features.src.paddings = paddings features.src.weights = weights features.src.task_ids = tf.cast(features.src.weights, dtype=tf.int32) * src_lang_ids features.src.ids_indicator = weights features.tgt = py_utils.NestedMap() features.tgt.ids = mass_out.tgt.ids features.tgt.labels = mass_out.tgt.labels features.tgt.paddings = paddings features.tgt.weights = mass_out.tgt.weights features.tgt.task_ids = tf.ones_like(features.src.task_ids, dtype=tf.int32) * tgt_lang_ids features.tgt.ids_indicator = weights if not py_utils.use_tpu(): features.src.strs = src features.tgt.strs = src return features.Transform(tf.squeeze)