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)
