def ComputePredictions(self, theta, input_batch):
"""Computes predictions for `input_batch`.
Args:
theta: A `.NestedMap` object containing variable values of this task.
input_batch: A `.NestedMap` expected to contain lasers.points_xyz,
lasers.points_feature, lasers.points_padding, cell_center_xyz,
cell_points_xyz, cell_feature, anchor_bboxes,
anchor_localization_residuals, assigned_gt_labels, and
assigned_cls_mask. See class doc string for details.
Returns:
A `.NestedMap` object containing residuals and classification_logits.
"""
p = self.params
input_batch.Transform(lambda x:
(x.shape, x.shape.num_elements())).VLog(
1, 'input_batch shapes: ')
cell_feature = py_utils.HasRank(input_batch.cell_feature, 4)
batch_size, num_centers = py_utils.GetShape(cell_feature, 2)
featurized_cell = self._CellFeaturizer(theta, input_batch)
# Project each featurized_cell features to each bbox per center.
featurized_anchors = self.cell_feature_projector.FProp(
theta.cell_feature_projector, featurized_cell)
# Reshape output so that we have features per offset.
featurized_anchors = tf.reshape(
featurized_anchors,
[batch_size, num_centers, p.num_anchor_bboxes_offsets, -1])
# Predict localization residuals.
predicted_residuals = self.localization_regressor.FProp(
theta.localization_regressor, featurized_anchors)
predicted_residuals = tf.reshape(
predicted_residuals,
[batch_size, num_centers, p.num_anchor_bboxes_per_center, 7])
if any([p.oracle_location, p.oracle_dimension, p.oracle_rotation]):
gt_residuals = py_utils.HasShape(
input_batch.anchor_localization_residuals,
[batch_size, num_centers, p.num_anchor_bboxes_per_center, 7])
residuals = []
if p.oracle_location:
residuals.append(gt_residuals[..., 0:3])
else:
residuals.append(predicted_residuals[..., 0:3])
if p.oracle_dimension:
residuals.append(gt_residuals[..., 3:6])
else:
residuals.append(predicted_residuals[..., 3:6])
if p.oracle_rotation:
residuals.append(gt_residuals[..., 6:])
else:
residuals.append(predicted_residuals[..., 6:])
predicted_residuals = tf.concat(residuals, axis=-1)
if p.squash_rotation_predictions:
predicted_rotations = predicted_residuals[..., 6:]
predicted_rotations = np.pi * tf.tanh(predicted_rotations)
predicted_residuals = tf.concat(
[predicted_residuals[..., :6], predicted_rotations], axis=-1)
# Predict object classification at each bbox.
predicted_classification_logits = self.classifier.FProp(
theta.classifier, featurized_anchors)
predicted_classification_logits = tf.reshape(
predicted_classification_logits, [
batch_size, num_centers, p.num_anchor_bboxes_per_center,
p.num_classes
])
if p.oracle_classification:
assigned_gt_labels = py_utils.HasShape(
input_batch.assigned_gt_labels,
[batch_size, num_centers, p.num_anchor_bboxes_per_center])
predicted_classification_logits = tf.one_hot(
assigned_gt_labels, p.num_classes)
return py_utils.NestedMap({
'residuals':
predicted_residuals,
'classification_logits':
predicted_classification_logits,
})
def ComputePredictions(self, theta, input_batch):
"""Computes predictions for `input_batch`.
Args:
theta: A `.NestedMap` object containing variable values of this task.
input_batch: A `.NestedMap` expected to contain cell_center_xyz,
cell_points_xyz, cell_feature, anchor_bboxes,
anchor_localization_residuals, assigned_gt_labels, and
assigned_cls_mask. See class doc string for details.
Returns:
A `.NestedMap` object containing residuals and classification_logits.
"""
p = self.params
input_batch.Transform(lambda x:
(x.shape, x.shape.num_elements())).VLog(
1, 'input_batch shapes: ')
cell_feature = py_utils.HasRank(input_batch.cell_feature, 4)
batch_size, num_centers, num_points_per_cell = py_utils.GetShape(
cell_feature, 3)
cell_points_xyz = py_utils.HasShape(
input_batch.cell_points_xyz,
[batch_size, num_centers, num_points_per_cell, 3])
cell_center_xyz = py_utils.HasShape(input_batch.cell_center_xyz,
[batch_size, num_centers, 3])
cell_points_padding = py_utils.HasShape(
input_batch.cell_points_padding,
[batch_size, num_centers, num_points_per_cell])
# TODO(jngiam): Make concat_feature computation a layer or configureable.
cell_center_xyz = tf.reshape(cell_center_xyz,
[batch_size, num_centers, 1, 3])
centered_cell_points_xyz = cell_points_xyz - cell_center_xyz
concat_feature = tf.concat([
tf.tile(cell_center_xyz, [1, 1, num_points_per_cell, 1]),
centered_cell_points_xyz, cell_feature
],
axis=-1) # pyformat: disable
# Featurize point clouds at each center.
point_input = py_utils.NestedMap({
'points': centered_cell_points_xyz,
'features': concat_feature,
'padding': cell_points_padding,
})
featurized_cell = self.cell_featurizer.FProp(theta.cell_featurizer,
point_input)
featurized_cell = py_utils.HasShape(featurized_cell,
[batch_size, num_centers, -1])
# Predict localization residuals.
predicted_residuals = self.localization_regressor.FProp(
theta.localization_regressor, featurized_cell)
predicted_residuals = tf.reshape(
predicted_residuals,
[batch_size, num_centers, p.num_anchor_bboxes_per_center, 7])
if p.squash_rotation_predictions:
predicted_rotations = predicted_residuals[..., 6:]
predicted_rotations = np.pi * tf.tanh(predicted_rotations)
predicted_residuals = tf.concat(
[predicted_residuals[..., :6], predicted_rotations], axis=-1)
# Predict object classification at each bbox.
predicted_classification_logits = self.classifier.FProp(
theta.classifier, featurized_cell)
predicted_classification_logits = tf.reshape(
predicted_classification_logits, [
batch_size, num_centers, p.num_anchor_bboxes_per_center,
p.num_classes
])
return py_utils.NestedMap({
'residuals':
predicted_residuals,
'classification_logits':
predicted_classification_logits,
})
def _GatedTanhFn(inputs): gated_inputs, act_inputs = tf.split(inputs, 2, axis=-1) return tf.tanh(act_inputs) * tf.sigmoid(gated_inputs)
def ComputePredictions(self, theta, input_batch):
"""Computes predictions for `input_batch`.
Args:
theta: A `.NestedMap` object containing variable values of this task.
input_batch: A `.NestedMap` object containing input tensors to this tower.
Returns:
A `.NestedMap` contains
logits - [b, nx, ny, nz, na, 7 + num_classes]
"""
p = self.params
input_batch.Transform(lambda x:
(x.shape, x.shape.num_elements())).VLog(
0, 'input_batch shapes: ')
# Make pillars representation from input_batch.
dense_features = self.input_featurizer.FProp(theta.input_featurizer,
input_batch)
# Backbone
tf.logging.vlog(1, 'dense_features.shape = %s', dense_features.shape)
act = self.backbone.FProp(theta.backbone, dense_features)
tf.logging.vlog(1, 'act.shape = %s', act.shape)
# Convert the output of the backbone into class logits and regression
# residuals using two different layers.
class_detection = self.class_detector.FProp(theta.class_detector, act)
reg_detection = self.regression_detector.FProp(
theta.regression_detector, act)
bs, nx, ny, _ = py_utils.GetShape(class_detection, 4)
predicted_classification_logits = tf.reshape(
class_detection,
[bs, nx, ny, p.grid_size_z, p.num_anchors, p.num_classes])
predicted_residuals = tf.reshape(
reg_detection, [bs, nx, ny, p.grid_size_z, p.num_anchors, 7])
if p.squash_rotation_predictions:
predicted_rotations = predicted_residuals[..., 6:]
predicted_rotations = np.pi * tf.tanh(predicted_rotations)
predicted_residuals = tf.concat(
[predicted_residuals[..., :6], predicted_rotations], axis=-1)
if p.oracle_location or p.oracle_dimension or p.oracle_rotation:
gt_residuals = py_utils.HasShape(
input_batch.anchor_localization_residuals,
[bs, nx, ny, p.grid_size_z, p.num_anchors, 7])
# Replace the predicted components with the ground truth if needed.
if p.oracle_location:
location = gt_residuals[..., 0:3]
else:
location = predicted_residuals[..., 0:3]
if p.oracle_dimension:
dimension = gt_residuals[..., 3:6]
else:
dimension = predicted_residuals[..., 3:6]
if p.oracle_rotation:
rotation = gt_residuals[..., 6:]
else:
rotation = predicted_residuals[..., 6:]
predicted_residuals = tf.concat([location, dimension, rotation],
axis=-1)
ret = py_utils.NestedMap({
'residuals':
predicted_residuals,
'classification_logits':
predicted_classification_logits,
})
if p.direction_classifier_weight > 0.0:
predicted_dir = self.direction_classifier.FProp(
theta.direction_classifier, act)
predicted_dir = tf.reshape(
predicted_dir, [bs, nx, ny, p.grid_size_z, p.num_anchors, 2])
ret.predicted_dir = predicted_dir
return ret
tf.nn.relu,
'RELU6':
tf.nn.relu6,
'LEAKY_RELU':
tf.nn.leaky_relu,
'SIGMOID':
tf.sigmoid,
'TANH':
tf.tanh,
'GELU':
tf.nn.gelu,
'GELU_APPROXIMATE':
lambda x: tf.nn.gelu(x, approximate=True),
'GELU_RAW':
lambda x: 0.5 * x * ( # pylint: disable=g-long-lambda
1 + tf.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))),
'SWISH':
tf.nn.swish,
'SOFTPLUS':
tf.nn.softplus,
# Squared ReLU from the Primer paper: https://arxiv.org/abs/2109.08668
'SQUARED_RELU':
lambda x: tf.math.square(tf.nn.relu(x)),
'SILU':
tf.nn.silu,
'NONE':
tf.identity,
}
_ACTIVATIONS_FLOPS = {
'NONE': 0,
def ComputePredictions(self, theta, input_batch):
"""Computes predictions for `input_batch`.
Args:
theta: A `.NestedMap` object containing variable values of this task.
input_batch: A `.NestedMap` object containing input tensors to this tower.
Returns:
A `.NestedMap` contains
logits - [b, nx, ny, nz, na, 7 + num_classes]
"""
p = self.params
input_batch.Transform(lambda x:
(x.shape, x.shape.num_elements())).VLog(
0, 'input_batch shapes: ')
bs, nx, ny, nz = py_utils.GetShape(input_batch.grid_num_points, 4)
# Process points to concatenate a set of fixed features (e.g.,
# add means, centers, normalize points to means).
num_features = 3 + p.num_laser_features
pillar_points = py_utils.HasShape(input_batch.pillar_points,
[bs, -1, -1, num_features])
_, npillars, npoints, _ = py_utils.GetShape(pillar_points, 4)
pillar_xyz = pillar_points[..., :3]
pillar_means = tf.reduce_mean(pillar_xyz, axis=2, keep_dims=True)
pillar_feats = pillar_points[..., 3:]
pillar_centers = py_utils.HasShape(input_batch.pillar_centers,
[bs, -1, 1, 3])
pillar_concat = tf.concat(axis=3,
values=[
pillar_xyz - pillar_means, pillar_feats,
tf.tile(pillar_means,
[1, 1, npoints, 1]),
tf.tile(pillar_centers,
[1, 1, npoints, 1])
])
# Featurize pillars.
pillar_features = self.featurizer.FProp(theta.featurizer,
pillar_concat)
# Convert back to the dense grid.
pillar_locations = py_utils.HasShape(input_batch.pillar_locations,
[bs, npillars, 3])
dense_features = _SparseToDense(grid_shape=(nx, ny, nz),
locations=pillar_locations,
feats=pillar_features)
# Backbone
tf.logging.vlog(1, 'dense_features.shape = %s', dense_features.shape)
act = self.backbone.FProp(theta.backbone, dense_features)
tf.logging.vlog(1, 'act.shape = %s', act.shape)
# Convert the output of the backbone into class logits and regression
# residuals using two different layers.
class_detection = self.class_detector.FProp(theta.class_detector, act)
reg_detection = self.regression_detector.FProp(
theta.regression_detector, act)
bs, nx, ny, _ = py_utils.GetShape(class_detection, 4)
predicted_classification_logits = tf.reshape(
class_detection,
[bs, nx, ny, p.grid_size_z, p.num_anchors, p.num_classes])
predicted_residuals = tf.reshape(
reg_detection, [bs, nx, ny, p.grid_size_z, p.num_anchors, 7])
if p.squash_rotation_predictions:
predicted_rotations = predicted_residuals[..., 6:]
predicted_rotations = np.pi * tf.tanh(predicted_rotations)
predicted_residuals = tf.concat(
[predicted_residuals[..., :6], predicted_rotations], axis=-1)
if p.oracle_location or p.oracle_dimension or p.oracle_rotation:
gt_residuals = py_utils.HasShape(
input_batch.anchor_localization_residuals,
[bs, nx, ny, p.grid_size_z, p.num_anchors, 7])
# Replace the predicted components with the ground truth if needed.
if p.oracle_location:
location = gt_residuals[..., 0:3]
else:
location = predicted_residuals[..., 0:3]
if p.oracle_dimension:
dimension = gt_residuals[..., 3:6]
else:
dimension = predicted_residuals[..., 3:6]
if p.oracle_rotation:
rotation = gt_residuals[..., 6:]
else:
rotation = predicted_residuals[..., 6:]
predicted_residuals = tf.concat([location, dimension, rotation],
axis=-1)
ret = py_utils.NestedMap({
'residuals':
predicted_residuals,
'classification_logits':
predicted_classification_logits,
})
if p.direction_classifier_weight > 0.0:
predicted_dir = self.direction_classifier.FProp(
theta.direction_classifier, act)
predicted_dir = tf.reshape(
predicted_dir, [bs, nx, ny, p.grid_size_z, p.num_anchors, 2])
ret.predicted_dir = predicted_dir
return ret
