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