def fwd_fn(query_queries_fwd, query_values_fwd, support_keys_fwd,
support_values_fwd, labels_fwd):
"""CrossTransformer forward, using a while loop to save memory."""
initial = (0,
tf.zeros([tf.reduce_max(labels) + 1, zero_dim],
dtype=query_queries_fwd.dtype))
def loop_body(idx, dist):
dist_new = self._get_dist(query_queries_fwd[idx:idx + 1],
query_values_fwd[idx:idx + 1],
support_keys_fwd, support_values_fwd,
labels_fwd)
dist = tf.concat([dist, dist_new], axis=1)
return (idx + 1, dist)
_, res = tf.while_loop(
lambda x, _: x < tf.shape(query_queries_fwd)[0],
loop_body,
initial,
parallel_iterations=1)
return res
def _voxel_hard_negative_classification_loss_unbatched(inputs_1, outputs_1,
is_intermediate, gamma):
"""Loss function for input and outputs of batch size 1."""
inputs_1, outputs_1 = _get_voxels_valid_inputs_outputs(inputs_1=inputs_1,
outputs_1=outputs_1)
if is_intermediate:
logits = outputs_1[standard_fields.DetectionResultFields.
intermediate_object_semantic_voxels]
else:
logits = outputs_1[
standard_fields.DetectionResultFields.object_semantic_voxels]
labels = tf.reshape(
inputs_1[standard_fields.InputDataFields.object_class_voxels], [-1])
background_mask = tf.equal(labels, 0)
num_background_points = tf.reduce_sum(
tf.cast(background_mask, dtype=tf.int32))
def loss_fn():
"""Loss function."""
num_classes = logits.get_shape().as_list()[-1]
if num_classes is None:
raise ValueError('Number of classes is unknown.')
masked_logits = tf.boolean_mask(logits, background_mask)
masked_weights = tf.pow(
1.0 - tf.reshape(tf.nn.softmax(masked_logits)[:, 0], [-1, 1]),
gamma)
num_points = tf.shape(masked_logits)[0]
masked_weights = masked_weights * tf.cast(
num_points, dtype=tf.float32) / tf.reduce_sum(masked_weights)
masked_labels_one_hot = tf.one_hot(indices=tf.boolean_mask(
labels, background_mask),
depth=num_classes)
loss = classification_loss_fn(logits=masked_logits,
labels=masked_labels_one_hot,
weights=masked_weights)
return loss
cond = tf.logical_and(tf.greater(num_background_points, 0),
tf.greater(tf.shape(labels)[0], 0))
return tf.cond(cond, loss_fn, lambda: tf.constant(0.0, dtype=tf.float32))
def identity_knn_graph(points, num_valid_points, k): # pylint: disable=unused-argument
"""Returns each points as its own neighbor k times.
Args:
points: A tf.float32 tensor of size [num_batches, N, D] where D is the point
dimensions.
num_valid_points: A tf.int32 tensor of size [num_batches] containing the
number of valid points in each batch example.
k: Number of neighbors for each point.
Returns:
distances: A tf.float32 tensor of [num_batches, N, k]. Distances is all
zeros since each point is returned as its own neighbor.
indices: A tf.int32 tensor of [num_batches, N, k]. Each row will contain
values that are identical to the index of that row.
"""
num_batches = points.get_shape()[0]
num_points = tf.shape(points)[1]
indices = tf.expand_dims(tf.range(num_points), axis=1)
indices = tf.tile(tf.expand_dims(indices, axis=0), [num_batches, 1, k])
distances = tf.zeros([num_batches, num_points, k], dtype=tf.float32)
return distances, indices
def sparse_voxel_grid_to_pointcloud(voxel_features, segment_ids,
num_valid_voxels, num_valid_points):
"""Convert voxel features back to points given their segment ids.
Args:
voxel_features: A tf.float32 tensor of size [batch_size, N', F].
segment_ids: A size [batch_size, N] tf.int32 tensor of IDs for each point
indicating which (flattened) voxel cell its data was mapped to.
num_valid_voxels: A tf.int32 tensor of size [batch_size] containing the
number of valid voxels in each batch example.
num_valid_points: A tf.int32 tensor of size [batch_size] containing the
number of valid points in each batch example.
Returns:
point_features: A tf.float32 tensor of size [batch_size, N, F].
Raises:
ValueError: If batch_size is unknown at graph construction time.
"""
batch_size = voxel_features.shape[0]
if batch_size is None:
raise ValueError('batch_size is unknown at graph construction time.')
num_points = tf.shape(segment_ids)[1]
def fn(i):
num_valid_voxels_i = num_valid_voxels[i]
num_valid_points_i = num_valid_points[i]
voxel_features_i = voxel_features[i, :num_valid_voxels_i, :]
segment_ids_i = segment_ids[i, :num_valid_points_i]
point_features = tf.gather(voxel_features_i, segment_ids_i)
point_features_rank = len(point_features.get_shape().as_list())
point_features_paddings = [[0, num_points - num_valid_points_i]]
for _ in range(point_features_rank - 1):
point_features_paddings.append([0, 0])
point_features = tf.pad(point_features,
paddings=point_features_paddings)
return point_features
return tf.map_fn(fn=fn, elems=tf.range(batch_size), dtype=tf.float32)
def fn():
"""Loss function for when number of input and output boxes is positive."""
if is_balanced:
weights = loss_utils.get_balanced_loss_weights_multiclass(
labels=input_boxes_instance_id)
else:
weights = tf.ones([tf.shape(input_boxes_instance_id)[0], 1],
dtype=tf.float32)
gt_rotation_matrix = tf.reshape(input_boxes_rotation_matrix, [-1, 9])
predicted_rotation_matrix = tf.reshape(output_boxes_rotation_matrix,
[-1, 9])
if loss_type == 'huber':
loss_fn = tf.keras.losses.Huber(
delta=delta, reduction=tf.keras.losses.Reduction.NONE)
elif loss_type == 'absolute_difference':
loss_fn = tf.keras.losses.MeanAbsoluteError(
reduction=tf.keras.losses.Reduction.NONE)
else:
raise ValueError(('Unknown loss type %s.' % loss_type))
rotation_losses = loss_fn(
y_true=gt_rotation_matrix, y_pred=predicted_rotation_matrix)
return tf.reduce_mean(rotation_losses * tf.reshape(weights, [-1]))
def knn_graph_from_points_unbatched(points, k, distance_upper_bound, mask=None):
"""Returns the distances and indices of the neighbors of each point.
Note that each point will have at least k neighbors unless the number of
points is less than k. In that case, the python function that is wrapped in
py_function will raise a value error.
Args:
points: A tf.float32 tensor of size [N, D] where D is the point dimensions.
k: Number of neighbors for each point.
distance_upper_bound: Only build the graph using points that are closer than
this distance.
mask: If not None, A tf.bool tensor of size [N]. If None, it is ignored.
If not None, knn will be applied to only points where the mask is True.
The points where the mask is False will have themselves as their
neighbors.
Returns:
distances: A tf.float32 tensor of size [N, k].
indices: A tf.int32 tensor of [N, k].
"""
def fn(np_points, np_mask):
return np_knn_graph_from_points_unbatched(
points=np_points,
k=k,
distance_upper_bound=distance_upper_bound,
mask=np_mask)
num_points = tf.shape(points)[0]
if mask is None:
mask = tf.cast(tf.ones([num_points], dtype=tf.int32), dtype=tf.bool)
else:
mask = tf.reshape(mask, [num_points])
distances, indices = tf.compat.v1.py_func(fn, [points, mask],
[tf.float32, tf.int32])
distances = tf.reshape(distances, [num_points, k])
indices = tf.reshape(indices, [num_points, k])
return distances, indices
def fn(i):
"""Map function."""
num_valid_points_i = num_valid_points[i]
points_i = points[i, :num_valid_points_i, :]
features_i = features[i, :num_valid_points_i, :]
voxel_features_i, voxel_indices_i, segment_ids_i, voxel_start_location_i = (
pointcloud_to_sparse_voxel_grid_unbatched(
points=points_i,
features=features_i,
grid_cell_size=grid_cell_size,
segment_func=segment_func))
num_valid_voxels_i = tf.shape(voxel_features_i)[0]
(voxel_features_i, voxel_indices_i, num_valid_voxels_i,
segment_ids_i) = _pad_or_clip_voxels(
voxel_features=voxel_features_i,
voxel_indices=voxel_indices_i,
num_valid_voxels=num_valid_voxels_i,
segment_ids=segment_ids_i,
voxels_pad_or_clip_size=voxels_pad_or_clip_size)
segment_ids_i = tf.pad(
segment_ids_i, paddings=[[0, num_points - num_valid_points_i]])
return (voxel_features_i, voxel_indices_i, num_valid_voxels_i,
segment_ids_i, voxel_start_location_i)
def select_slate_greedy(slate_size, s_no_click, s, q):
"""Selects the slate using the adaptive greedy algorithm.
This algorithm corresponds to the method "GS" in
Ie et al. https://arxiv.org/abs/1905.12767.
Args:
slate_size: int, the size of the recommendation slate.
s_no_click: float tensor, the score for not clicking any document.
s: [num_of_documents] tensor, the scores for clicking documents.
q: [num_of_documents] tensor, the predicted q values for documents.
Returns:
[slate_size] tensor, the selected slate.
"""
def argmax(v, mask):
return tf.argmax(
input=(v - tf.reduce_min(input_tensor=v) + 1) * mask, axis=0)
numerator = tf.constant(0.)
denominator = tf.constant(0.) + s_no_click
mask = tf.ones(tf.shape(input=q)[0])
def set_element(v, i, x):
mask = tf.one_hot(i, tf.shape(input=v)[0])
v_new = tf.ones_like(v) * x
return tf.where(tf.equal(mask, 1), v_new, v)
for _ in range(slate_size):
k = argmax((numerator + s * q) / (denominator + s), mask)
mask = set_element(mask, k, 0)
numerator = numerator + tf.gather(s * q, k)
denominator = denominator + tf.gather(s, k)
output_slate = tf.where(tf.equal(mask, 0))
return output_slate
def _build_train_op(self):
"""Builds the training op for Rainbow.
Returns:
train_op: An op performing one step of training.
"""
target_distribution = tf.stop_gradient(
self._build_target_distribution())
# size of indices: batch_size x 1.
indices = tf.range(tf.shape(self._replay_logits)[0])[:, None]
# size of reshaped_actions: batch_size x 2.
reshaped_actions = tf.concat([indices, self._replay.actions[:, None]],
1)
# For each element of the batch, fetch the logits for its selected action.
chosen_action_logits = tf.gather_nd(self._replay_logits,
reshaped_actions)
loss = tf.nn.softmax_cross_entropy_with_logits(
labels=target_distribution, logits=chosen_action_logits)
optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate,
epsilon=self.optimizer_epsilon)
update_priorities_op = self._replay.tf_set_priority(
self._replay.indices, tf.sqrt(loss + 1e-10))
target_priorities = self._replay.tf_get_priority(self._replay.indices)
target_priorities = tf.math.add(target_priorities, 1e-10)
target_priorities = 1.0 / tf.sqrt(target_priorities)
target_priorities /= tf.reduce_max(target_priorities)
weighted_loss = target_priorities * loss
with tf.control_dependencies([update_priorities_op]):
return optimizer.minimize(
tf.reduce_mean(weighted_loss)), weighted_loss
def decode_with_label(self, example_string):
"""Processes a single example string.
Extracts and processes the image, and ignores the label. We assume that the
image has three channels.
Args:
example_string: str, an Example protocol buffer.
Returns:
image_rescaled: the image, resized to `image_size x image_size` and
rescaled to [-1, 1]. Note that Gaussian data augmentation may cause
values to go beyond this range.
label: tf.int
"""
ex_decoded = read_example_and_parse_image(example_string)
image_decoded = ex_decoded['image']
image_resized = tf.image.resize_images(
image_decoded, [self.image_size, self.image_size],
method=tf.image.ResizeMethod.BILINEAR,
align_corners=True)
image_resized = tf.cast(image_resized, tf.float32)
image = 2 * (image_resized / 255.0 - 0.5) # Rescale to [-1, 1].
if self.data_augmentation is not None:
if self.data_augmentation.enable_gaussian_noise:
image = image + tf.random_normal(tf.shape(
image)) * self.data_augmentation.gaussian_noise_std
if self.data_augmentation.enable_jitter:
j = self.data_augmentation.jitter_amount
paddings = tf.constant([[j, j], [j, j], [0, 0]])
image = tf.pad(image, paddings, 'REFLECT')
image = tf.image.random_crop(
image, [self.image_size, self.image_size, 3])
return image, tf.cast(ex_decoded['label'], dtype=tf.int32)
def set_element(v, i, x):
mask = tf.one_hot(i, tf.shape(input=v)[0])
v_new = tf.ones_like(v) * x
return tf.where(tf.equal(mask, 1), v_new, v)
def classification_loss_using_mask_iou(inputs,
outputs,
num_samples,
max_instance_id=None,
similarity_strategy='distance',
is_balanced=True,
is_intermediate=False):
"""Classification loss with an iou threshold.
Args:
inputs: A dictionary that contains
num_valid_voxels - A tf.int32 tensor of size [batch_size].
instance_ids - A tf.int32 tensor of size [batch_size, n].
class_labels - A tf.int32 tensor of size [batch_size, n]. It is assumed
that the background voxels are assigned to class 0.
outputs: A dictionart that contains
embeddings - A tf.float32 tensor of size [batch_size, n, f].
logits - A tf.float32 tensor of size [batch_size, n, num_classes]. It is
assumed that background is class 0.
num_samples: An int determining the number of samples.
max_instance_id: If set, instance ids larger than that value will be
ignored. If not set, it will be computed from instance_ids tensor.
similarity_strategy: Defines the method for computing similarity between
embedding vectors. Possible values are 'dotproduct' and
'distance'.
is_balanced: If True, the per-voxel losses are re-weighted to have equal
total weight for foreground vs. background voxels.
is_intermediate: True if applied to intermediate predictions;
otherwise, False.
Returns:
A tf.float32 scalar loss tensor.
"""
instance_ids_key = standard_fields.InputDataFields.object_instance_id_voxels
class_labels_key = standard_fields.InputDataFields.object_class_voxels
num_voxels_key = standard_fields.InputDataFields.num_valid_voxels
if is_intermediate:
embedding_key = (standard_fields.DetectionResultFields.
intermediate_instance_embedding_voxels)
logits_key = (standard_fields.DetectionResultFields.
intermediate_object_semantic_voxels)
else:
embedding_key = (
standard_fields.DetectionResultFields.instance_embedding_voxels)
logits_key = standard_fields.DetectionResultFields.object_semantic_voxels
if instance_ids_key not in inputs:
raise ValueError('instance_ids is missing in inputs.')
if class_labels_key not in inputs:
raise ValueError('class_labels is missing in inputs.')
if num_voxels_key not in inputs:
raise ValueError('num_voxels is missing in inputs.')
if embedding_key not in outputs:
raise ValueError('embedding is missing in outputs.')
if logits_key not in outputs:
raise ValueError('logits is missing in outputs.')
batch_size = inputs[num_voxels_key].get_shape().as_list()[0]
if batch_size is None:
raise ValueError(
'batch_size is not defined at graph construction time.')
num_valid_voxels = inputs[num_voxels_key]
num_voxels = tf.shape(inputs[instance_ids_key])[1]
valid_mask = tf.less(
tf.tile(tf.expand_dims(tf.range(num_voxels), axis=0), [batch_size, 1]),
tf.expand_dims(num_valid_voxels, axis=1))
return classification_loss_using_mask_iou_func(
embeddings=outputs[embedding_key],
logits=outputs[logits_key],
instance_ids=tf.reshape(inputs[instance_ids_key], [batch_size, -1]),
class_labels=inputs[class_labels_key],
num_samples=num_samples,
valid_mask=valid_mask,
max_instance_id=max_instance_id,
similarity_strategy=similarity_strategy,
is_balanced=is_balanced)
def _box_corner_distance_loss(
loss_type, is_balanced, input_boxes_length, input_boxes_height,
input_boxes_width, input_boxes_center, input_boxes_rotation_matrix,
input_boxes_instance_id, output_boxes_length, output_boxes_height,
output_boxes_width, output_boxes_center, output_boxes_rotation_matrix,
delta):
"""Computes regression loss on object corner locations."""
def fn():
"""Loss function for when number of input and output boxes is positive."""
if is_balanced:
weights = loss_utils.get_balanced_loss_weights_multiclass(
labels=input_boxes_instance_id)
else:
weights = tf.ones([tf.shape(input_boxes_instance_id)[0], 1],
dtype=tf.float32)
normalized_box_size = 5.0
predicted_boxes_length = output_boxes_length
predicted_boxes_height = output_boxes_height
predicted_boxes_width = output_boxes_width
predicted_boxes_center = output_boxes_center
predicted_boxes_rotation_matrix = output_boxes_rotation_matrix
gt_boxes_length = input_boxes_length
gt_boxes_height = input_boxes_height
gt_boxes_width = input_boxes_width
gt_boxes_center = input_boxes_center
gt_boxes_rotation_matrix = input_boxes_rotation_matrix
if loss_type in ['normalized_huber', 'normalized_euclidean']:
predicted_boxes_length /= (gt_boxes_length / normalized_box_size)
predicted_boxes_height /= (gt_boxes_height / normalized_box_size)
predicted_boxes_width /= (gt_boxes_width / normalized_box_size)
gt_boxes_length = tf.ones_like(
gt_boxes_length, dtype=tf.float32) * normalized_box_size
gt_boxes_height = tf.ones_like(
gt_boxes_height, dtype=tf.float32) * normalized_box_size
gt_boxes_width = tf.ones_like(
gt_boxes_width, dtype=tf.float32) * normalized_box_size
gt_box_corners = box_utils.get_box_corners_3d(
boxes_length=gt_boxes_length,
boxes_height=gt_boxes_height,
boxes_width=gt_boxes_width,
boxes_rotation_matrix=gt_boxes_rotation_matrix,
boxes_center=gt_boxes_center)
predicted_box_corners = box_utils.get_box_corners_3d(
boxes_length=predicted_boxes_length,
boxes_height=predicted_boxes_height,
boxes_width=predicted_boxes_width,
boxes_rotation_matrix=predicted_boxes_rotation_matrix,
boxes_center=predicted_boxes_center)
corner_weights = tf.tile(weights, [1, 8])
if loss_type in ['huber', 'normalized_huber']:
loss_fn = tf.keras.losses.Huber(
delta=delta, reduction=tf.keras.losses.Reduction.NONE)
elif loss_type in ['normalized_absolute_difference', 'absolute_difference']:
loss_fn = tf.keras.losses.MeanAbsoluteError(
reduction=tf.keras.losses.Reduction.NONE)
else:
raise ValueError(('Unknown loss type %s.' % loss_type))
box_corner_losses = loss_fn(
y_true=tf.reshape(gt_box_corners, [-1, 3]),
y_pred=tf.reshape(predicted_box_corners, [-1, 3]))
return tf.reduce_mean(box_corner_losses * tf.reshape(corner_weights, [-1]))
cond_input = tf.greater(tf.shape(input_boxes_length)[0], 0)
cond_output = tf.greater(tf.shape(output_boxes_length)[0], 0)
cond = tf.logical_and(cond_input, cond_output)
return tf.cond(cond, fn, lambda: tf.constant(0.0, dtype=tf.float32))
def compute_motion_labels(scene,
frame0,
frame1,
frame_start_index,
points_key,
box_margin=0.1):
"""Compute motion label for each point.
Args:
scene: dict of tensor containing scene.
frame0: dict of tensor containing points and objects.
frame1: dict of tensor containing points and objects.
frame_start_index: starting frame index.
points_key: A string corresponding to the tensor of point positions in
inputs.
box_margin: A margin value to enlarge box, so that surrounding points are
included.
Returns:
A motion tensor of [N, 3] shape.
"""
point_positions = frame0[points_key]
frame0_object_names = frame0['objects/name']
frame1_object_names = frame1['objects/name']
bool_matrix = tf.math.equal(
tf.expand_dims(frame0_object_names, axis=1),
tf.expand_dims(frame1_object_names, axis=0))
match_indices = tf.where(bool_matrix)
# object box level
box_dimension = tf.gather(
frame0['objects/shape/dimension'], match_indices[:, 0], axis=0)
boxes_length = box_dimension[:, 0:1]
boxes_width = box_dimension[:, 1:2]
boxes_height = box_dimension[:, 2:3]
boxes_rotation_matrix = tf.gather(
frame0['objects/pose/R'], match_indices[:, 0], axis=0)
boxes_center = tf.gather(
frame0['objects/pose/t'], match_indices[:, 0], axis=0)
frame1_box_rotation_matrix = tf.gather(
frame1['objects/pose/R'], match_indices[:, 1], axis=0)
frame1_box_center = tf.gather(
frame1['objects/pose/t'], match_indices[:, 1], axis=0)
# frame level
frame0_rotation = scene['frames/pose/R'][frame_start_index]
frame1_rotation = scene['frames/pose/R'][frame_start_index + 1]
frame0_translation = scene['frames/pose/t'][frame_start_index]
frame1_translation = scene['frames/pose/t'][frame_start_index + 1]
frame1_box_center_global = tf.tensordot(
frame1_box_center, frame1_rotation, axes=(1, 1)) + frame1_translation
frame1_box_center_in_frame0 = tf.tensordot(
frame1_box_center_global - frame0_translation,
frame0_rotation,
axes=(1, 0))
# only find index on boxes that are matched between two frames
points_box_index = box_utils.map_points_to_boxes(
points=point_positions,
boxes_length=boxes_length,
boxes_height=boxes_height,
boxes_width=boxes_width,
boxes_rotation_matrix=boxes_rotation_matrix,
boxes_center=boxes_center,
box_margin=box_margin)
# TODO(huangrui): disappered object box have 0 motion.
# Probably consider set to nan or ignore_label.
# 1. gather points in surviving matched box only,
# and replicate rotation/t to same length;
# 2. get points in box frame, apply new rotation/t per point;
# 3. new location minus old location -> motion vector;
# 4. scatter it to a larger motion_vector with 0 for
# points ouside of matched boxes.
# Need to limit boxes to those matched boxes.
# otherwise the points_box_index will contain useless box.
# index in all point array, of points that are inside the box.
points_inside_box_index = tf.where(points_box_index + 1)[:, 0]
box_index = tf.gather(points_box_index, points_inside_box_index)
points_inside_box = tf.gather(point_positions, points_inside_box_index)
box_rotation_per_point = tf.gather(boxes_rotation_matrix, box_index)
box_center_per_point = tf.gather(boxes_center, box_index)
# Tensor [N, 3, 3] and [N, 3]. note we are transform points reversely.
points_in_box_frame = tf.einsum('ikj,ik->ij', box_rotation_per_point,
points_inside_box - box_center_per_point)
# Transform rotation of box from frame1 coordinate to frame0 coordinate
# note, transpose is implemented via changing summation axis
frame1_box_rotation_matrix_global = tf.transpose(
tf.tensordot(frame1_rotation, frame1_box_rotation_matrix, axes=(1, 1)),
perm=(1, 0, 2))
frame1_box_rotation_matrix_in_frame0 = tf.transpose(
tf.tensordot(
frame0_rotation, frame1_box_rotation_matrix_global, axes=(0, 1)),
perm=(1, 0, 2))
# this is the points_position_after_following_frame1_box's motion.
frame1_box_rotation_in_frame0_per_point = tf.gather(
frame1_box_rotation_matrix_in_frame0, box_index)
frame1_box_center_in_frame0_per_point = tf.gather(frame1_box_center_in_frame0,
box_index)
points_in_box_frame1 = tf.einsum(
'ijk,ik->ij', frame1_box_rotation_in_frame0_per_point,
points_in_box_frame) + frame1_box_center_in_frame0_per_point
motion_vector = points_in_box_frame1 - points_inside_box
scattered_vector = tf.scatter_nd(
indices=tf.expand_dims(points_inside_box_index, axis=1),
updates=motion_vector,
shape=tf.shape(point_positions, out_type=tf.dtypes.int64))
return scattered_vector
def _box_classification_using_center_distance_loss_unbatched(
inputs_1, outputs_1, is_intermediate, is_balanced,
max_positive_normalized_distance):
"""Loss function for input and outputs of batch size 1."""
inputs_1, outputs_1 = _get_voxels_valid_inputs_outputs(inputs_1=inputs_1,
outputs_1=outputs_1)
if is_intermediate:
output_object_centers = outputs_1[standard_fields.DetectionResultFields
.intermediate_object_center_voxels]
output_object_length = outputs_1[standard_fields.DetectionResultFields.
intermediate_object_length_voxels]
output_object_height = outputs_1[standard_fields.DetectionResultFields.
intermediate_object_height_voxels]
output_object_width = outputs_1[standard_fields.DetectionResultFields.
intermediate_object_width_voxels]
output_object_rotation_matrix = outputs_1[
standard_fields.DetectionResultFields.
intermediate_object_rotation_matrix_voxels]
logits = outputs_1[standard_fields.DetectionResultFields.
intermediate_object_semantic_voxels]
else:
output_object_centers = outputs_1[
standard_fields.DetectionResultFields.object_center_voxels]
output_object_length = outputs_1[
standard_fields.DetectionResultFields.object_length_voxels]
output_object_height = outputs_1[
standard_fields.DetectionResultFields.object_height_voxels]
output_object_width = outputs_1[
standard_fields.DetectionResultFields.object_width_voxels]
output_object_rotation_matrix = outputs_1[
standard_fields.DetectionResultFields.
object_rotation_matrix_voxels]
logits = outputs_1[
standard_fields.DetectionResultFields.object_semantic_voxels]
normalized_center_distance = loss_utils.get_normalized_corner_distances(
predicted_boxes_center=output_object_centers,
predicted_boxes_length=output_object_length,
predicted_boxes_height=output_object_height,
predicted_boxes_width=output_object_width,
predicted_boxes_rotation_matrix=output_object_rotation_matrix,
gt_boxes_center=inputs_1[
standard_fields.InputDataFields.object_center_voxels],
gt_boxes_length=inputs_1[
standard_fields.InputDataFields.object_length_voxels],
gt_boxes_height=inputs_1[
standard_fields.InputDataFields.object_height_voxels],
gt_boxes_width=inputs_1[
standard_fields.InputDataFields.object_width_voxels],
gt_boxes_rotation_matrix=inputs_1[
standard_fields.InputDataFields.object_rotation_matrix_voxels])
labels = tf.reshape(
inputs_1[standard_fields.InputDataFields.object_class_voxels], [-1])
instances = tf.reshape(
inputs_1[standard_fields.InputDataFields.object_instance_id_voxels],
[-1])
params = {}
if is_balanced:
weights = loss_utils.get_balanced_loss_weights_multiclass(
labels=tf.expand_dims(instances, axis=1))
params['weights'] = weights
def loss_fn():
"""Loss function."""
num_classes = logits.get_shape().as_list()[-1]
if num_classes is None:
raise ValueError('Number of classes is unknown.')
labels_one_hot = tf.one_hot(indices=(labels - 1),
depth=(num_classes - 1))
inverse_distance_coef = tf.maximum(
tf.minimum(
1.0 -
normalized_center_distance / max_positive_normalized_distance,
1.0), 0.0)
labels_one_hot = tf.reshape(inverse_distance_coef,
[-1, 1]) * labels_one_hot
background_label = 1.0 - tf.math.reduce_sum(
labels_one_hot, axis=1, keepdims=True)
labels_one_hot = tf.concat([background_label, labels_one_hot], axis=1)
loss = classification_loss_fn(logits=logits,
labels=labels_one_hot,
**params)
return loss
return tf.cond(tf.greater(tf.shape(labels)[0], 0), loss_fn,
lambda: tf.constant(0.0, dtype=tf.float32))
def pointcloud_to_sparse_voxel_grid(points, features, num_valid_points,
grid_cell_size, voxels_pad_or_clip_size,
segment_func):
"""Converts a pointcloud into a voxel grid.
This function calls the `pointcloud_to_sparse_voxel_grid_unbatched`
function avove in a while loop to map a batch of points to a batch of voxels.
Args:
points: A tf.float32 tensor of size [batch_size, N, 3].
features: A tf.float32 tensor of size [batch_size, N, F].
num_valid_points: A tf.int32 tensor of size [num_batches] containing the
number of valid points in each batch example.
grid_cell_size: A tf.float32 tensor of size [3].
voxels_pad_or_clip_size: Number of target voxels to pad or clip to. If None,
it will not perform the padding.
segment_func: A tensorflow function that operates on segments. Examples are
one of tf.math.unsorted_segment_{min/max/mean/prod/sum}.
Returns:
voxel_features: A tf.float32 tensor of size [batch_size, N', F]
or [batch_size, N', G, F] where G is the number of points sampled per
voxel.
voxel_indices: A tf.int32 tensor of size [batch_size, N', 3].
num_valid_voxels: A tf.int32 tensor of size [batch_size].
segment_ids: A size [batch_size, N] tf.int32 tensor of IDs for each point
indicating which (flattened) voxel cell its data was mapped to.
voxel_start_location: A size [batch_size, 3] tf.float32 tensor of voxel
start locations.
Raises:
ValueError: If pooling method is unknown.
"""
batch_size = points.get_shape().as_list()[0]
if batch_size is None:
batch_size = tf.shape(points)[0]
num_points = tf.shape(points)[1]
def fn(i):
"""Map function."""
num_valid_points_i = num_valid_points[i]
points_i = points[i, :num_valid_points_i, :]
features_i = features[i, :num_valid_points_i, :]
voxel_features_i, voxel_indices_i, segment_ids_i, voxel_start_location_i = (
pointcloud_to_sparse_voxel_grid_unbatched(
points=points_i,
features=features_i,
grid_cell_size=grid_cell_size,
segment_func=segment_func))
num_valid_voxels_i = tf.shape(voxel_features_i)[0]
(voxel_features_i, voxel_indices_i, num_valid_voxels_i,
segment_ids_i) = _pad_or_clip_voxels(
voxel_features=voxel_features_i,
voxel_indices=voxel_indices_i,
num_valid_voxels=num_valid_voxels_i,
segment_ids=segment_ids_i,
voxels_pad_or_clip_size=voxels_pad_or_clip_size)
segment_ids_i = tf.pad(
segment_ids_i, paddings=[[0, num_points - num_valid_points_i]])
return (voxel_features_i, voxel_indices_i, num_valid_voxels_i,
segment_ids_i, voxel_start_location_i)
return tf.map_fn(
fn=fn,
elems=tf.range(batch_size),
dtype=(tf.float32, tf.int32, tf.int32, tf.int32, tf.float32))
def _build_train_op(self, optimizer):
"""Build the TensorFlow graph used to learn the bisimulation metric.
Args:
optimizer: a tf.train optimizer.
Returns:
A TensorFlow op to minimize the bisimulation loss.
"""
self.online_network = tf.make_template('Online',
self._network_template)
self.target_network = tf.make_template('Target',
self._network_template)
self.s1_ph = tf.placeholder(tf.float64, (self.batch_size, 2),
name='s1_ph')
self.s2_ph = tf.placeholder(tf.float64, (self.batch_size, 2),
name='s2_ph')
self.s1_online_distances = self.online_network(
self._concat_states(self.s1_ph))
self.s1_target_distances = self.target_network(
self._concat_states(self.s1_ph))
self.s2_target_distances = self.target_network(
self._concat_states(self.s2_ph))
self.action_ph = tf.placeholder(tf.int32, (self.batch_size,))
self.rewards_ph = tf.placeholder(tf.float64, (self.batch_size,))
# We use an expanding horizon for computing the distances.
self.bisim_horizon_ph = tf.placeholder(tf.float64, ())
# bisimulation_target_1 = rew_diff + gamma * next_distance.
bisimulation_target_1 = tf.stop_gradient(self._build_bisimulation_target())
# bisimulation_target_2 = curr_distance.
bisimulation_target_2 = tf.stop_gradient(self.s1_target_distances)
# We slowly taper in the maximum according to the bisim horizon.
bisimulation_target = tf.maximum(
bisimulation_target_1, bisimulation_target_2 * self.bisim_horizon_ph)
# We zero-out diagonal entries, since those are estimating the distance
# between a state and itself, which we know to be 0.
diagonal_mask = 1.0 - tf.diag(tf.ones(self.batch_size, dtype=tf.float64))
diagonal_mask = tf.reshape(diagonal_mask, (self.batch_size**2, 1))
bisimulation_target *= diagonal_mask
bisimulation_estimate = self.s1_online_distances
# We start with a mask that includes everything.
loss_mask = tf.ones(tf.shape(bisimulation_estimate))
# We have to enforce that states being compared are done only using the same
# action.
indicators = self.action_ph
indicators = tf.cast(indicators, tf.float64)
# indicators will initially have shape [batch_size], we first tile it:
square_ids = tf.tile([indicators], [self.batch_size, 1])
# We subtract square_ids from its transpose:
square_ids = square_ids - tf.transpose(square_ids)
# At this point all zero-entries are the ones with equal IDs.
# Now we would like to convert the zeros in this matrix to 1s, and make
# everything else a 0. We can do this with the following operation:
loss_mask = 1 - tf.abs(tf.sign(square_ids))
# Now reshape to match the shapes of the estimate and target.
loss_mask = tf.reshape(loss_mask, (self.batch_size**2, 1))
larger_targets = bisimulation_target - bisimulation_estimate
larger_targets_count = tf.reduce_sum(
tf.cast(larger_targets > 0., tf.float64))
tf.summary.scalar('Learning/LargerTargets', larger_targets_count)
tf.summary.scalar('Learning/NumUpdates', tf.count_nonzero(loss_mask))
tf.summary.scalar('Learning/BisimHorizon', self.bisim_horizon_ph)
bisimulation_loss = tf.losses.mean_squared_error(
bisimulation_target,
bisimulation_estimate,
weights=loss_mask)
tf.summary.scalar('Learning/loss', bisimulation_loss)
# Plot average distance between sampled representations.
average_distance = tf.reduce_mean(bisimulation_estimate)
tf.summary.scalar('Approx/AverageDistance', average_distance)
return optimizer.minimize(bisimulation_loss)
def random_eraser(images,
min_size,
max_size,
probability,
max_operations,
probability_additional_operations,
augment_entire_batch = False):
"""Earses a random rectangle shaped areas in the second image or image batch.
Args:
images: Stacked image pair that should be augmented with shape
[2, height, width, 3] or a batch of images that should be augmented with
shape [batch, height, width, 3].
min_size: Minimum size of erased rectangle.
max_size: Maximum size of erased rectangle.
probability: Probability of applying this augementation function.
max_operations: Maximum number total areas that should be erased.
probability_additional_operations: Probability for each additional area to
be erased if augementation is applied.
augment_entire_batch: If true the input is treated as batch of images to
which the augmentation should be applid.
Returns:
Possibly augemented images.
"""
perform_erase = tf.less(tf.random.uniform([]), probability)
height = tf.shape(images)[-3]
width = tf.shape(images)[-2]
# Returns augemented images.
def true_fn(images):
if augment_entire_batch:
image_2 = images
mean_color = tf.reduce_mean(image_2, axis=[1, 2], keepdims=True)
print(mean_color.shape)
else:
image_1, image_2 = tf.unstack(images)
mean_color = tf.reduce_mean(image_2, axis=[0, 1], keepdims=True)
def body(var_img, mean_color):
x0 = tf.random.uniform([], 0, width, dtype=tf.int32)
y0 = tf.random.uniform([], 0, height, dtype=tf.int32)
dx = tf.random.uniform([], min_size, max_size, dtype=tf.int32)
dy = tf.random.uniform([], min_size, max_size, dtype=tf.int32)
x = tf.range(width)
x_mask = (x0 <= x) & (x < x0+dx)
y = tf.range(height)
y_mask = (y0 <= y) & (y < y0+dy)
mask = x_mask & y_mask[:, tf.newaxis]
mask = tf.cast(mask[:, :, tf.newaxis], image_2.dtype)
result = var_img * (1 - mask) + mean_color * mask
return result
# Perform at least one erase operation.
image_2 = body(image_2, mean_color)
# Perform additional erase operations.
for _ in range(max_operations - 1):
perform_erase = tf.less(
tf.random.uniform([]), probability_additional_operations)
image_2 = tf.cond(perform_erase, lambda: body(image_2, mean_color),
lambda: image_2)
if augment_entire_batch:
images = image_2
else:
images = tf.stack([image_1, image_2])
return images
# Returns unaugmented images.
def false_fn(images):
return images
return tf.cond(perform_erase,
lambda: true_fn(images),
lambda: false_fn(images))
def random_crop(images,
flow,
mask,
crop_height,
crop_width,
relative_offset,
probability_crop_offset):
"""Performs a random crop with the given height and width."""
# early return if crop_height or crop_width is not specified
if crop_height is None or crop_width is None:
return images, flow, mask
orig_height = tf.shape(images)[-3]
orig_width = tf.shape(images)[-2]
# check if crop size fits the image size
scale = 1.0
ratio = tf.cast(crop_height, tf.float32) / tf.cast(orig_height, tf.float32)
scale = tf.math.maximum(scale, ratio)
ratio = tf.cast(crop_width, tf.float32) / tf.cast(orig_width, tf.float32)
scale = tf.math.maximum(scale, ratio)
# compute minimum required hight
new_height = tf.cast(
tf.math.ceil(tf.cast(orig_height, tf.float32) * scale), tf.int32)
new_width = tf.cast(
tf.math.ceil(tf.cast(orig_width, tf.float32) * scale), tf.int32)
# perform resize (scales with 1 if not required)
images = smurf_utils.resize(images, new_height, new_width, is_flow=False)
# compute joint offset
max_offset_h = new_height - tf.cast(crop_height, dtype=tf.int32)
max_offset_w = new_width - tf.cast(crop_width, dtype=tf.int32)
joint_offset_h = tf.random.uniform([], maxval=max_offset_h+1, dtype=tf.int32)
joint_offset_w = tf.random.uniform([], maxval=max_offset_w+1, dtype=tf.int32)
# compute relative offset
min_relative_offset_h = tf.math.maximum(
joint_offset_h - relative_offset, 0)
max_relative_offset_h = tf.math.minimum(
joint_offset_h + relative_offset, max_offset_h)
min_relative_offset_w = tf.math.maximum(
joint_offset_w - relative_offset, 0)
max_relative_offset_w = tf.math.minimum(
joint_offset_w + relative_offset, max_offset_w)
relative_offset_h = tf.random.uniform(
[], minval=min_relative_offset_h, maxval=max_relative_offset_h+1,
dtype=tf.int32)
relative_offset_w = tf.random.uniform(
[], minval=min_relative_offset_w, maxval=max_relative_offset_w+1,
dtype=tf.int32)
set_crop_offset = tf.random.uniform([]) < probability_crop_offset
relative_offset_h = tf.cond(
set_crop_offset, lambda: relative_offset_h, lambda: joint_offset_h)
relative_offset_w = tf.cond(
set_crop_offset, lambda: relative_offset_w, lambda: joint_offset_w)
# crop both images
image_1, image_2 = tf.unstack(images)
image_1 = tf.image.crop_to_bounding_box(
image_1, offset_height=joint_offset_h, offset_width=joint_offset_w,
target_height=crop_height, target_width=crop_width)
image_2 = tf.image.crop_to_bounding_box(
image_2, offset_height=relative_offset_h, offset_width=relative_offset_w,
target_height=crop_height, target_width=crop_width)
images = tf.stack([image_1, image_2])
if flow is not None:
# perform resize (scales with 1 if not required)
flow, mask = smurf_utils.resize(
flow, new_height, new_width, is_flow=True, mask=mask)
# crop flow and mask
flow = tf.image.crop_to_bounding_box(
flow,
offset_height=joint_offset_h,
offset_width=joint_offset_w,
target_height=crop_height,
target_width=crop_width)
mask = tf.image.crop_to_bounding_box(
mask,
offset_height=joint_offset_h,
offset_width=joint_offset_w,
target_height=crop_height,
target_width=crop_width)
# correct flow for relative shift (/crop)
flow_delta = tf.stack(
[tf.cast(relative_offset_h - joint_offset_h, tf.float32),
tf.cast(relative_offset_w - joint_offset_w, tf.float32)])
flow = (flow - flow_delta) * mask
return images, flow, mask, joint_offset_h, joint_offset_w
def project_distribution(supports, weights, target_support,
validate_args=False):
"""Projects a batch of (support, weights) onto target_support.
Based on equation (7) in (Bellemare et al., 2017):
https://arxiv.org/abs/1707.06887
In the rest of the comments we will refer to this equation simply as Eq7.
This code is not easy to digest, so we will use a running example to clarify
what is going on, with the following sample inputs:
* supports = [[0, 2, 4, 6, 8],
[1, 3, 4, 5, 6]]
* weights = [[0.1, 0.6, 0.1, 0.1, 0.1],
[0.1, 0.2, 0.5, 0.1, 0.1]]
* target_support = [4, 5, 6, 7, 8]
In the code below, comments preceded with 'Ex:' will be referencing the above
values.
Args:
supports: Tensor of shape (batch_size, num_dims) defining supports for the
distribution.
weights: Tensor of shape (batch_size, num_dims) defining weights on the
original support points. Although for the CategoricalDQN agent these
weights are probabilities, it is not required that they are.
target_support: Tensor of shape (num_dims) defining support of the projected
distribution. The values must be monotonically increasing. Vmin and Vmax
will be inferred from the first and last elements of this tensor,
respectively. The values in this tensor must be equally spaced.
validate_args: Whether we will verify the contents of the
target_support parameter.
Returns:
A Tensor of shape (batch_size, num_dims) with the projection of a batch of
(support, weights) onto target_support.
Raises:
ValueError: If target_support has no dimensions, or if shapes of supports,
weights, and target_support are incompatible.
"""
target_support_deltas = target_support[1:] - target_support[:-1]
# delta_z = `\Delta z` in Eq7.
delta_z = target_support_deltas[0]
validate_deps = []
supports.shape.assert_is_compatible_with(weights.shape)
supports[0].shape.assert_is_compatible_with(target_support.shape)
target_support.shape.assert_has_rank(1)
if validate_args:
# Assert that supports and weights have the same shapes.
validate_deps.append(
tf.Assert(
tf.reduce_all(tf.equal(tf.shape(supports), tf.shape(weights))),
[supports, weights]))
# Assert that elements of supports and target_support have the same shape.
validate_deps.append(
tf.Assert(
tf.reduce_all(
tf.equal(tf.shape(supports)[1], tf.shape(target_support))),
[supports, target_support]))
# Assert that target_support has a single dimension.
validate_deps.append(
tf.Assert(
tf.equal(tf.size(tf.shape(target_support)), 1), [target_support]))
# Assert that the target_support is monotonically increasing.
validate_deps.append(
tf.Assert(tf.reduce_all(target_support_deltas > 0), [target_support]))
# Assert that the values in target_support are equally spaced.
validate_deps.append(
tf.Assert(
tf.reduce_all(tf.equal(target_support_deltas, delta_z)),
[target_support]))
with tf.control_dependencies(validate_deps):
# Ex: `v_min, v_max = 4, 8`.
v_min, v_max = target_support[0], target_support[-1]
# Ex: `batch_size = 2`.
batch_size = tf.shape(supports)[0]
# `N` in Eq7.
# Ex: `num_dims = 5`.
num_dims = tf.shape(target_support)[0]
# clipped_support = `[\hat{T}_{z_j}]^{V_max}_{V_min}` in Eq7.
# Ex: `clipped_support = [[[ 4. 4. 4. 6. 8.]]
# [[ 4. 4. 4. 5. 6.]]]`.
clipped_support = tf.clip_by_value(supports, v_min, v_max)[:, None, :]
# Ex: `tiled_support = [[[[ 4. 4. 4. 6. 8.]
# [ 4. 4. 4. 6. 8.]
# [ 4. 4. 4. 6. 8.]
# [ 4. 4. 4. 6. 8.]
# [ 4. 4. 4. 6. 8.]]
# [[ 4. 4. 4. 5. 6.]
# [ 4. 4. 4. 5. 6.]
# [ 4. 4. 4. 5. 6.]
# [ 4. 4. 4. 5. 6.]
# [ 4. 4. 4. 5. 6.]]]]`.
tiled_support = tf.tile([clipped_support], [1, 1, num_dims, 1])
# Ex: `reshaped_target_support = [[[ 4.]
# [ 5.]
# [ 6.]
# [ 7.]
# [ 8.]]
# [[ 4.]
# [ 5.]
# [ 6.]
# [ 7.]
# [ 8.]]]`.
reshaped_target_support = tf.tile(target_support[:, None], [batch_size, 1])
reshaped_target_support = tf.reshape(reshaped_target_support,
[batch_size, num_dims, 1])
# numerator = `|clipped_support - z_i|` in Eq7.
# Ex: `numerator = [[[[ 0. 0. 0. 2. 4.]
# [ 1. 1. 1. 1. 3.]
# [ 2. 2. 2. 0. 2.]
# [ 3. 3. 3. 1. 1.]
# [ 4. 4. 4. 2. 0.]]
# [[ 0. 0. 0. 1. 2.]
# [ 1. 1. 1. 0. 1.]
# [ 2. 2. 2. 1. 0.]
# [ 3. 3. 3. 2. 1.]
# [ 4. 4. 4. 3. 2.]]]]`.
numerator = tf.abs(tiled_support - reshaped_target_support)
quotient = 1 - (numerator / delta_z)
# clipped_quotient = `[1 - numerator / (\Delta z)]_0^1` in Eq7.
# Ex: `clipped_quotient = [[[[ 1. 1. 1. 0. 0.]
# [ 0. 0. 0. 0. 0.]
# [ 0. 0. 0. 1. 0.]
# [ 0. 0. 0. 0. 0.]
# [ 0. 0. 0. 0. 1.]]
# [[ 1. 1. 1. 0. 0.]
# [ 0. 0. 0. 1. 0.]
# [ 0. 0. 0. 0. 1.]
# [ 0. 0. 0. 0. 0.]
# [ 0. 0. 0. 0. 0.]]]]`.
clipped_quotient = tf.clip_by_value(quotient, 0, 1)
# Ex: `weights = [[ 0.1 0.6 0.1 0.1 0.1]
# [ 0.1 0.2 0.5 0.1 0.1]]`.
weights = weights[:, None, :]
# inner_prod = `\sum_{j=0}^{N-1} clipped_quotient * p_j(x', \pi(x'))`
# in Eq7.
# Ex: `inner_prod = [[[[ 0.1 0.6 0.1 0. 0. ]
# [ 0. 0. 0. 0. 0. ]
# [ 0. 0. 0. 0.1 0. ]
# [ 0. 0. 0. 0. 0. ]
# [ 0. 0. 0. 0. 0.1]]
# [[ 0.1 0.2 0.5 0. 0. ]
# [ 0. 0. 0. 0.1 0. ]
# [ 0. 0. 0. 0. 0.1]
# [ 0. 0. 0. 0. 0. ]
# [ 0. 0. 0. 0. 0. ]]]]`.
inner_prod = clipped_quotient * weights
# Ex: `projection = [[ 0.8 0.0 0.1 0.0 0.1]
# [ 0.8 0.1 0.1 0.0 0.0]]`.
projection = tf.reduce_sum(inner_prod, 3)
projection = tf.reshape(projection, [batch_size, num_dims])
return projection
def knn_graph_from_points(points,
num_valid_points,
k,
distance_upper_bound,
mask=None):
"""Returns the distances and indices of the neighbors of each point.
Note that each point will have at least k neighbors unless the number of
points is less than k. In that case, the python function that is wrapped in
py_function will raise a value error.
Args:
points: A tf.float32 tensor of size [batch_size, N, D] where D is the point
dimensions.
num_valid_points: A tf.int32 tensor of size [batch_size] containing the
number of valid points in each batch example.
k: Number of neighbors for each point.
distance_upper_bound: Only build the graph using points that are closer than
this distance.
mask: If not None, A tf.bool tensor of size [batch_size, N]. If None, it is
ignored. If not None, knn will be applied to only points where the mask is
True. The points where the mask is False will have themselves as their
neighbors.
Returns:
distances: A tf.float32 tensor of size [batch_size, N, k].
indices: A tf.int32 tensor of size [batch_size, N, k].
Raises:
ValueError: If batch_size is unknown.
"""
if points.get_shape().as_list()[0] is None:
raise ValueError('Batch size is unknown.')
batch_size = points.get_shape().as_list()[0]
num_points = tf.shape(points)[1]
def fn_knn_graph_from_points_unbatched(i):
"""Computes knn graph for example i in the batch."""
num_valid_points_i = num_valid_points[i]
points_i = points[i, :num_valid_points_i, :]
if mask is None:
mask_i = None
else:
mask_i = mask[i, :num_valid_points_i]
distances_i, indices_i = knn_graph_from_points_unbatched(
points=points_i,
k=k,
distance_upper_bound=distance_upper_bound,
mask=mask_i)
distances_i = tf.pad(distances_i,
paddings=[[0, num_points - num_valid_points_i],
[0, 0]])
indices_i = tf.pad(indices_i,
paddings=[[0, num_points - num_valid_points_i],
[0, 0]])
return distances_i, indices_i
distances, indices = tf.map_fn(fn=fn_knn_graph_from_points_unbatched,
elems=tf.range(batch_size),
dtype=(tf.float32, tf.int32))
return distances, indices
def preprocess(inputs,
output_keys=None,
is_training=False,
using_sequence_dataset=False,
num_frame_to_load=1,
transform_points_fn=None,
image_preprocess_fn_dic=None,
images_points_correspondence_fn=None,
compute_semantic_labels_fn=None,
compute_motion_labels_fn=None,
view_names=(),
points_key='points',
colors_key='colors',
normals_key='normals',
intensities_key='intensities',
elongations_key='elongations',
semantic_labels_key='semantic_labels',
motion_labels_key='motion_labels',
spin_coords_key=None,
points_in_image_frame_key=None,
num_points_to_randomly_sample=None,
x_min_degree_rotation=None,
x_max_degree_rotation=None,
y_min_degree_rotation=None,
y_max_degree_rotation=None,
z_min_degree_rotation=None,
z_max_degree_rotation=None,
points_pad_or_clip_size=None,
voxels_pad_or_clip_size=None,
voxel_grid_cell_size=(0.1, 0.1, 0.1),
num_offset_bins_x=4,
num_offset_bins_y=4,
num_offset_bins_z=4,
point_feature_keys=('point_offsets', ),
point_to_voxel_segment_func=tf.math.unsorted_segment_mean,
x_random_crop_size=None,
y_random_crop_size=None,
min_scale_ratio=None,
max_scale_ratio=None,
semantic_labels_offset=0,
ignore_labels=(),
remove_unlabeled_images_and_points=False,
labeled_view_name=None,
only_keep_first_return_lidar_points=False):
"""Preprocesses a dictionary of `Tensor` inputs.
If is_training=True, it will randomly rotate the points around the z axis,
and will randomly flip the points with respect to x and/or y axis.
Note that the preprocessor function does not correct normal vectors if they
exist in the inputs.
Note that the preprocessing effects all values of `inputs` that are `Tensors`.
Args:
inputs: A dictionary of inputs. Each value must be a `Tensor`.
output_keys: Either None, or a list of strings containing the keys in the
dictionary that is returned by the preprocess function.
is_training: Whether we're training or testing.
using_sequence_dataset: if true, the inputs will contain scene and multiple
frames data.
num_frame_to_load: If greater than 1, load multiframe point cloud point
positions and its correspondence.
transform_points_fn: Fn to transform other frames to a specific frame's
coordinate.
image_preprocess_fn_dic: Image preprocessing function. Maps view names to
their image preprocessing functions. Set it to None, if there are no
images to preprocess or you are not interested in preprocesing images.
images_points_correspondence_fn: The function that computes correspondence
between images and points.
compute_semantic_labels_fn: If not None, semantic labels will be computed
using this function.
compute_motion_labels_fn: If not None, motion labels will be computed using
this function.
view_names: Names corresponding to 2d views of the scene.
points_key: The key used for `points` in the inputs.
colors_key: The key used for `colors` in the inputs.
normals_key: The key used for 'normals' in the inputs.
intensities_key: The key used for 'intensities' in the inputs.
elongations_key: The key used for 'elongations' in the inputs.
semantic_labels_key: The key used for 'semantic_labels' in the inputs.
motion_labels_key: The key used for 'motion_labels' in the inputs.
spin_coords_key: The key used for 'spin_coords' in the inputs. In Waymo
data, spin_coords is a [num_points, 3] tensor that contains scan_index,
shot_index, return_index. In Waymo data, return_index of the first return
points is 0.
points_in_image_frame_key: A string that identifies the tensor that contains
the points_in_image_frame tensor. If None, it won't be used.
num_points_to_randomly_sample: Number of points to randomly sample. If None,
it will keep the original points and does not perform sampling.
x_min_degree_rotation: Min degree of rotation around the x axis.
x_max_degree_rotation: Max degree of ratation around the x axis.
y_min_degree_rotation: Min degree of rotation around the y axis.
y_max_degree_rotation: Max degree of ratation around the y axis.
z_min_degree_rotation: Min degree of rotation around the z axis.
z_max_degree_rotation: Max degree of ratation around the z axis.
points_pad_or_clip_size: Number of target points to pad or clip to. If None,
it will not perform the point padding.
voxels_pad_or_clip_size: Number of target voxels to pad or clip to. If None,
it will not perform the voxel padding.
voxel_grid_cell_size: A three dimensional tuple determining the voxel grid
size.
num_offset_bins_x: Number of bins for point offsets in x direction.
num_offset_bins_y: Number of bins for point offsets in y direction.
num_offset_bins_z: Number of bins for point offsets in z direction.
point_feature_keys: The keys used to form the voxel features.
point_to_voxel_segment_func: The function used to aggregate the features
of the points that fall in the same voxel.
x_random_crop_size: Size of the random crop in x dimension. If None, random
crop will not take place on x dimension.
y_random_crop_size: Size of the random crop in y dimension. If None, random
crop will not take place on y dimension.
min_scale_ratio: Minimum scale ratio. Used for scaling point cloud.
max_scale_ratio: Maximum scale ratio. Used for scaling point cloud.
semantic_labels_offset: An integer offset that will be added to labels.
ignore_labels: A tuple containing labels that should be ignored when
computing the loss and metrics.
remove_unlabeled_images_and_points: If True, removes the images that are not
labeled and also removes the points that are associated with those images.
labeled_view_name: The name of the view that is labeled, otherwise None.
only_keep_first_return_lidar_points: If True, we only keep the first return
lidar points.
Returns:
The mean subtracted points with an optional rotation applied.
Raises:
ValueError: if `inputs` doesn't contain the points_key.
ValueError: if `points_in_image_frame` does not have rank 3.
"""
inputs = dict(inputs)
if using_sequence_dataset:
all_frame_inputs = inputs
scene = all_frame_inputs['scene']
frame1 = all_frame_inputs['frame1']
frame_start_index = all_frame_inputs['frame_start_index']
inputs = dict(
all_frame_inputs['frame0']
) # so that the following processing code can be unchanged.
# Initializing empty dictionary for mesh, image, indices_2d and non tensor
# inputs.
non_tensor_inputs = {}
view_image_inputs = {}
view_indices_2d_inputs = {}
mesh_inputs = {}
if image_preprocess_fn_dic is None:
image_preprocess_fn_dic = {}
# Convert all float64 to float32 and all int64 to int32.
for key in sorted(inputs):
if isinstance(inputs[key], tf.Tensor):
if inputs[key].dtype == tf.float64:
inputs[key] = tf.cast(inputs[key], dtype=tf.float32)
if inputs[key].dtype == tf.int64:
inputs[key] = tf.cast(inputs[key], dtype=tf.int32)
if points_key in inputs:
inputs[standard_fields.InputDataFields.
point_positions] = inputs[points_key]
if colors_key is not None and colors_key in inputs:
inputs[
standard_fields.InputDataFields.point_colors] = inputs[colors_key]
if normals_key is not None and normals_key in inputs:
inputs[standard_fields.InputDataFields.
point_normals] = inputs[normals_key]
if intensities_key is not None and intensities_key in inputs:
inputs[standard_fields.InputDataFields.
point_intensities] = inputs[intensities_key]
if elongations_key is not None and elongations_key in inputs:
inputs[standard_fields.InputDataFields.
point_elongations] = inputs[elongations_key]
if semantic_labels_key is not None and semantic_labels_key in inputs:
inputs[standard_fields.InputDataFields.
object_class_points] = inputs[semantic_labels_key]
if motion_labels_key is not None and motion_labels_key in inputs:
inputs[standard_fields.InputDataFields.
object_flow_points] = inputs[motion_labels_key]
if spin_coords_key is not None and spin_coords_key in inputs:
inputs[standard_fields.InputDataFields.
point_spin_coordinates] = inputs[spin_coords_key]
# Acquire point / image correspondences.
if images_points_correspondence_fn is not None:
fn_outputs = images_points_correspondence_fn(inputs)
if 'points_position' in fn_outputs:
inputs[standard_fields.InputDataFields.
point_positions] = fn_outputs['points_position']
if 'points_intensity' in fn_outputs and intensities_key is not None:
inputs[standard_fields.InputDataFields.
point_intensities] = fn_outputs['points_intensity']
if 'points_elongation' in fn_outputs and elongations_key is not None:
inputs[standard_fields.InputDataFields.
point_elongations] = fn_outputs['points_elongation']
if 'points_label' in fn_outputs and semantic_labels_key is not None:
inputs[standard_fields.InputDataFields.
object_class_points] = fn_outputs['points_label']
if 'view_images' in fn_outputs:
for key in sorted(fn_outputs['view_images']):
if len(fn_outputs['view_images'][key].shape) != 4:
raise ValueError(('%s image should have rank 4.' % key))
view_image_inputs = fn_outputs['view_images']
if 'view_indices_2d' in fn_outputs:
for key in sorted(fn_outputs['view_indices_2d']):
if len(fn_outputs['view_indices_2d'][key].shape) != 3:
raise ValueError(
('%s indices_2d should have rank 3.' % key))
view_indices_2d_inputs = fn_outputs['view_indices_2d']
else:
if points_in_image_frame_key is not None:
inputs['rgb_view/features'] = inputs['image']
inputs['rgb_view/indices_2d'] = inputs[points_in_image_frame_key]
if len(inputs['rgb_view/indices_2d'].shape) != 3:
raise ValueError('`points_in_image_frame` should have rank 3.')
frame0 = inputs.copy()
if num_frame_to_load > 1:
point_positions_list = [
frame0[standard_fields.InputDataFields.point_positions]
]
if view_indices_2d_inputs:
view_indices_2d_list = [view_indices_2d_inputs[view_names[0]]]
frame_source_list = [
tf.zeros([
tf.shape(
frame0[standard_fields.InputDataFields.point_positions])[0]
], tf.int32)
]
for i in range(1, num_frame_to_load):
target_frame_key = 'frame' + str(i)
if images_points_correspondence_fn is not None:
frame_i = images_points_correspondence_fn(
all_frame_inputs[target_frame_key])
else:
raise ValueError(
'images_points_correspondence_fn is needed for loading multi-frame pointclouds.'
)
transformed_point_positions = transform_points_fn(
scene, frame_i['points_position'], frame_start_index,
i + frame_start_index)
point_positions_list.append(transformed_point_positions)
if view_indices_2d_inputs:
view_indices_2d_list.append(
frame_i['view_indices_2d'][view_names[0]])
frame_source_list.append(
tf.ones([tf.shape(transformed_point_positions)[0]], tf.int32) *
i)
# add multi-frame info to override inputs and view_indices_2d_inputs
inputs[standard_fields.InputDataFields.
point_frame_index] = tf.expand_dims(tf.concat(frame_source_list,
axis=0),
axis=1)
inputs[standard_fields.InputDataFields.point_positions] = tf.concat(
point_positions_list, axis=0)
if view_indices_2d_inputs:
view_indices_2d_inputs[view_names[0]] = tf.concat(
view_indices_2d_list, axis=1)
# Validate inputs.
if standard_fields.InputDataFields.point_positions not in inputs:
raise ValueError('`inputs` must contain a point_positions')
if inputs[
standard_fields.InputDataFields.point_positions].shape.ndims != 2:
raise ValueError('points must be of rank 2.')
if inputs[standard_fields.InputDataFields.point_positions].shape[1] != 3:
raise ValueError('point should be 3 dimensional.')
# Remove normal nans.
if standard_fields.InputDataFields.point_normals in inputs:
inputs[standard_fields.InputDataFields.point_normals] = tf.where(
tf.math.is_nan(
inputs[standard_fields.InputDataFields.point_normals]),
tf.zeros_like(
inputs[standard_fields.InputDataFields.point_normals]),
inputs[standard_fields.InputDataFields.point_normals])
# Compute semantic labels if compute_semantic_labels_fn is not None
# An example is when the ground-truth contains 3d object boxes and not per
# point labels. This would be a function that infers point labels from boxes.
if compute_semantic_labels_fn is not None:
inputs[standard_fields.InputDataFields.
object_class_points] = compute_semantic_labels_fn(
inputs=frame0,
points_key=standard_fields.InputDataFields.point_positions)
if compute_motion_labels_fn is not None:
inputs[standard_fields.InputDataFields.
object_flow_points] = compute_motion_labels_fn(
scene=scene,
frame0=frame0,
frame1=frame1,
frame_start_index=frame_start_index,
points_key=standard_fields.InputDataFields.point_positions)
# Splitting inputs to {view_image_inputs,
# view_indices_2d_inputs,
# mesh_inputs,
# non_tensor_inputs}
mesh_keys = []
for key in [
standard_fields.InputDataFields.point_positions,
standard_fields.InputDataFields.point_colors,
standard_fields.InputDataFields.point_normals,
standard_fields.InputDataFields.point_intensities,
standard_fields.InputDataFields.point_elongations,
standard_fields.InputDataFields.object_class_points,
standard_fields.InputDataFields.point_spin_coordinates,
standard_fields.InputDataFields.object_flow_points,
standard_fields.InputDataFields.point_frame_index,
]:
if key is not None and key in inputs:
mesh_keys.append(key)
view_image_names = [('%s/features' % key) for key in view_names]
view_indices_2d_names = [('%s/indices_2d' % key) for key in view_names]
# Additional key collecting
for k, v in six.iteritems(inputs):
if k in view_image_names:
view_image_inputs[k] = v
elif k in view_indices_2d_names:
view_indices_2d_inputs[k] = v
elif k in mesh_keys:
if num_frame_to_load > 1:
pad_size = tf.shape(
inputs[standard_fields.InputDataFields.
point_positions])[0] - tf.shape(v)[0]
if k == standard_fields.InputDataFields.object_class_points:
pad_value = -1
else:
pad_value = 0
v = tf.pad(v, [[0, pad_size], [0, 0]],
constant_values=pad_value)
mesh_inputs[k] = v
else:
non_tensor_inputs[k] = v
# Remove points that are not in the lidar first return (optional)
if only_keep_first_return_lidar_points:
_remove_second_return_lidar_points(
mesh_inputs=mesh_inputs,
view_indices_2d_inputs=view_indices_2d_inputs)
# Randomly sample points
preprocessor_utils.randomly_sample_points(
mesh_inputs=mesh_inputs,
view_indices_2d_inputs=view_indices_2d_inputs,
target_num_points=num_points_to_randomly_sample)
# Add weights if it does not exist in inputs. The weight of the points with
# label in `ignore_labels` is set to 0. This helps the loss and metrics to
# ignore those labels.
use_weights = (
standard_fields.InputDataFields.object_class_points in mesh_inputs
or standard_fields.InputDataFields.object_flow_points in mesh_inputs)
if use_weights:
if num_frame_to_load > 1:
num_valid_points_frame0 = tf.shape(
frame0[standard_fields.InputDataFields.point_positions])[0]
num_additional_frame_points = tf.shape(
mesh_inputs[standard_fields.InputDataFields.
object_class_points])[0] - num_valid_points_frame0
weights = tf.concat([
tf.ones([num_valid_points_frame0, 1], tf.float32),
tf.zeros([num_additional_frame_points, 1], tf.float32)
],
axis=0)
else:
weights = tf.ones_like(mesh_inputs[
standard_fields.InputDataFields.object_class_points],
dtype=tf.float32)
if standard_fields.InputDataFields.object_class_points in mesh_inputs:
mesh_inputs[
standard_fields.InputDataFields.object_class_points] = tf.cast(
mesh_inputs[
standard_fields.InputDataFields.object_class_points],
dtype=tf.int32)
for ignore_label in ignore_labels:
weights *= tf.cast(tf.not_equal(
mesh_inputs[
standard_fields.InputDataFields.object_class_points],
ignore_label),
dtype=tf.float32)
mesh_inputs[
standard_fields.InputDataFields.point_loss_weights] = weights
mesh_inputs[standard_fields.InputDataFields.
object_class_points] += semantic_labels_offset
# We normalize the intensities and elongations to be in a smaller range.
if standard_fields.InputDataFields.point_intensities in mesh_inputs:
mesh_inputs[standard_fields.InputDataFields.
point_intensities] = change_intensity_range(
intensities=mesh_inputs[
standard_fields.InputDataFields.point_intensities])
if standard_fields.InputDataFields.point_elongations in mesh_inputs:
mesh_inputs[
standard_fields.InputDataFields.point_elongations] = (tf.cast(
mesh_inputs[standard_fields.InputDataFields.point_elongations],
dtype=tf.float32) * 2.0 / 255.0) - 1.0
# Random scale the points.
if min_scale_ratio is not None and max_scale_ratio is not None:
scale_ratio = tf.random.uniform([],
minval=min_scale_ratio,
maxval=max_scale_ratio,
dtype=tf.float32)
mesh_inputs[
standard_fields.InputDataFields.point_positions] *= scale_ratio
if standard_fields.InputDataFields.object_flow_points in mesh_inputs:
mesh_inputs[standard_fields.InputDataFields.
object_flow_points] *= scale_ratio
# Random crop the points.
randomly_crop_points(mesh_inputs=mesh_inputs,
view_indices_2d_inputs=view_indices_2d_inputs,
x_random_crop_size=x_random_crop_size,
y_random_crop_size=y_random_crop_size)
# If training, pick the best labeled image and points that project to it.
# In many datasets, only one image is labeled anyways.
if remove_unlabeled_images_and_points:
pick_labeled_image(mesh_inputs=mesh_inputs,
view_image_inputs=view_image_inputs,
view_indices_2d_inputs=view_indices_2d_inputs,
view_name=labeled_view_name)
# Process images.
preprocessor_utils.preprocess_images(
view_image_inputs=view_image_inputs,
view_indices_2d_inputs=view_indices_2d_inputs,
image_preprocess_fn_dic=image_preprocess_fn_dic,
is_training=is_training)
# Record the original points.
original_points = mesh_inputs[
standard_fields.InputDataFields.point_positions]
if standard_fields.InputDataFields.point_colors in mesh_inputs:
original_colors = mesh_inputs[
standard_fields.InputDataFields.point_colors]
if standard_fields.InputDataFields.point_normals in mesh_inputs:
original_normals = mesh_inputs[
standard_fields.InputDataFields.point_normals]
# Update feature visibility count.
if 'feature_visibility_count' in mesh_inputs:
mesh_inputs['feature_visibility_count'] = tf.maximum(
mesh_inputs['feature_visibility_count'], 1)
mesh_inputs['features'] /= tf.cast(
mesh_inputs['feature_visibility_count'], dtype=tf.float32)
# Subtract mean from points.
mean_points = tf.reduce_mean(
mesh_inputs[standard_fields.InputDataFields.point_positions], axis=0)
mesh_inputs[
standard_fields.InputDataFields.point_positions] -= tf.expand_dims(
mean_points, axis=0)
# Rotate points randomly.
if standard_fields.InputDataFields.point_normals in mesh_inputs:
normals = mesh_inputs[standard_fields.InputDataFields.point_normals]
else:
normals = None
if standard_fields.InputDataFields.object_flow_points in mesh_inputs:
motions = mesh_inputs[
standard_fields.InputDataFields.object_flow_points]
else:
motions = None
(mesh_inputs[standard_fields.InputDataFields.point_positions],
rotated_normals, rotated_motions) = rotate_randomly(
points=mesh_inputs[standard_fields.InputDataFields.point_positions],
normals=normals,
motions=motions,
x_min_degree_rotation=x_min_degree_rotation,
x_max_degree_rotation=x_max_degree_rotation,
y_min_degree_rotation=y_min_degree_rotation,
y_max_degree_rotation=y_max_degree_rotation,
z_min_degree_rotation=z_min_degree_rotation,
z_max_degree_rotation=z_max_degree_rotation)
# Random flipping in x and y directions.
(mesh_inputs[standard_fields.InputDataFields.point_positions],
flipped_normals,
flipped_motions) = flip_randomly_points_and_normals_motions(
points=mesh_inputs[standard_fields.InputDataFields.point_positions],
normals=rotated_normals,
motions=rotated_motions,
is_training=is_training)
if standard_fields.InputDataFields.point_normals in mesh_inputs:
mesh_inputs[
standard_fields.InputDataFields.point_normals] = flipped_normals
if standard_fields.InputDataFields.object_flow_points in mesh_inputs:
mesh_inputs[standard_fields.InputDataFields.
object_flow_points] = flipped_motions
# Normalize RGB to [-1.0, 1.0].
if standard_fields.InputDataFields.point_colors in mesh_inputs:
mesh_inputs[standard_fields.InputDataFields.point_colors] = tf.cast(
mesh_inputs[standard_fields.InputDataFields.point_colors],
dtype=tf.float32)
mesh_inputs[standard_fields.InputDataFields.point_colors] *= (2.0 /
255.0)
mesh_inputs[standard_fields.InputDataFields.point_colors] -= 1.0
# Add original points to mesh inputs.
mesh_inputs[standard_fields.InputDataFields.
point_positions_original] = original_points
if standard_fields.InputDataFields.point_colors in mesh_inputs:
mesh_inputs[standard_fields.InputDataFields.
point_colors_original] = original_colors
if standard_fields.InputDataFields.point_normals in mesh_inputs:
mesh_inputs[standard_fields.InputDataFields.
point_normals_original] = original_normals
# Pad or clip the point tensors.
pad_or_clip(mesh_inputs=mesh_inputs,
view_indices_2d_inputs=view_indices_2d_inputs,
pad_or_clip_size=points_pad_or_clip_size)
if num_frame_to_load > 1:
# Note: num_valid_points is the sum of 'num_points_per_fram' for now.
# num_points_per_frame is each frame's valid num of points.
# TODO(huangrui): if random sampling is called earlier, the count here
# is not guaranteed to be in order. need sorting.
if num_points_to_randomly_sample is not None:
raise ValueError(
'randomly sample is not compatible with padding multi frame point clouds yet!'
)
_, _, mesh_inputs[standard_fields.InputDataFields.
num_valid_points_per_frame] = tf.unique_with_counts(
tf.reshape(
mesh_inputs[standard_fields.InputDataFields.
point_frame_index], [-1]))
if points_pad_or_clip_size is not None:
padded_points = tf.where_v2(
tf.greater(
points_pad_or_clip_size, mesh_inputs[
standard_fields.InputDataFields.num_valid_points]),
points_pad_or_clip_size -
mesh_inputs[standard_fields.InputDataFields.num_valid_points],
0)
# Correct the potential unique count error from optionally padded 0s point
# frame index.
mesh_inputs[
standard_fields.InputDataFields.
num_valid_points_per_frame] -= tf.pad(
tf.expand_dims(padded_points, 0), [[
0,
tf.shape(mesh_inputs[standard_fields.InputDataFields.
num_valid_points_per_frame])[0] -
1
]])
# Putting back the dictionaries together
processed_inputs = mesh_inputs.copy()
processed_inputs.update(non_tensor_inputs)
for key in sorted(view_image_inputs):
processed_inputs[('%s/features' % key)] = view_image_inputs[key]
for key in sorted(view_indices_2d_inputs):
processed_inputs[('%s/indices_2d' % key)] = view_indices_2d_inputs[key]
# Create features that do not exist
if 'point_offsets' in point_feature_keys:
preprocessor_utils.add_point_offsets(
inputs=processed_inputs, voxel_grid_cell_size=voxel_grid_cell_size)
if 'point_offset_bins' in point_feature_keys:
preprocessor_utils.add_point_offset_bins(
inputs=processed_inputs,
voxel_grid_cell_size=voxel_grid_cell_size,
num_bins_x=num_offset_bins_x,
num_bins_y=num_offset_bins_y,
num_bins_z=num_offset_bins_z)
# Voxelize point features
preprocessor_utils.voxelize_point_features(
inputs=processed_inputs,
voxels_pad_or_clip_size=voxels_pad_or_clip_size,
voxel_grid_cell_size=voxel_grid_cell_size,
point_feature_keys=point_feature_keys,
point_to_voxel_segment_func=point_to_voxel_segment_func,
num_frame_to_load=num_frame_to_load)
# Voxelize point / image correspondence indices
preprocessor_utils.voxelize_point_to_view_correspondences(
inputs=processed_inputs,
view_indices_2d_inputs=view_indices_2d_inputs,
voxels_pad_or_clip_size=voxels_pad_or_clip_size,
voxel_grid_cell_size=voxel_grid_cell_size)
# Voxelizing the semantic labels
preprocessor_utils.voxelize_semantic_labels(
inputs=processed_inputs,
voxels_pad_or_clip_size=voxels_pad_or_clip_size,
voxel_grid_cell_size=voxel_grid_cell_size)
# Voxelizing the loss weights
preprocessor_utils.voxelize_property_tensor(
inputs=processed_inputs,
point_tensor_key=standard_fields.InputDataFields.point_loss_weights,
corresponding_voxel_tensor_key=standard_fields.InputDataFields.
voxel_loss_weights,
voxels_pad_or_clip_size=voxels_pad_or_clip_size,
voxel_grid_cell_size=voxel_grid_cell_size,
segment_func=tf.math.unsorted_segment_max)
# Voxelizing the object flow
if standard_fields.InputDataFields.object_flow_points in processed_inputs:
preprocessor_utils.voxelize_property_tensor(
inputs=processed_inputs,
point_tensor_key=standard_fields.InputDataFields.
object_flow_points,
corresponding_voxel_tensor_key='object_flow_voxels_max',
voxels_pad_or_clip_size=voxels_pad_or_clip_size,
voxel_grid_cell_size=voxel_grid_cell_size,
segment_func=tf.math.unsorted_segment_max)
preprocessor_utils.voxelize_property_tensor(
inputs=processed_inputs,
point_tensor_key=standard_fields.InputDataFields.
object_flow_points,
corresponding_voxel_tensor_key='object_flow_voxels_min',
voxels_pad_or_clip_size=voxels_pad_or_clip_size,
voxel_grid_cell_size=voxel_grid_cell_size,
segment_func=tf.math.unsorted_segment_min)
processed_inputs[standard_fields.InputDataFields.
object_flow_voxels] = processed_inputs[
'object_flow_voxels_max'] + processed_inputs[
'object_flow_voxels_min']
if num_frame_to_load > 1:
mesh_inputs[
standard_fields.InputDataFields.num_valid_points] = mesh_inputs[
standard_fields.InputDataFields.num_valid_points_per_frame][0]
# Filter preprocessed_inputs by output_keys if it is not None.
if output_keys is not None:
processed_inputs = {
k: v
for k, v in six.iteritems(processed_inputs) if k in output_keys
}
return processed_inputs
def randomly_crop_points(mesh_inputs,
view_indices_2d_inputs,
x_random_crop_size,
y_random_crop_size,
epsilon=1e-5):
"""Randomly crops points.
Args:
mesh_inputs: A dictionary containing input mesh (point) tensors.
view_indices_2d_inputs: A dictionary containing input point to view
correspondence tensors.
x_random_crop_size: Size of the random crop in x dimension. If None, random
crop will not take place on x dimension.
y_random_crop_size: Size of the random crop in y dimension. If None, random
crop will not take place on y dimension.
epsilon: Epsilon (a very small value) used to add as a small margin to
thresholds.
"""
if x_random_crop_size is None and y_random_crop_size is None:
return
points = mesh_inputs[standard_fields.InputDataFields.point_positions]
num_points = tf.shape(points)[0]
# Pick a random point
if x_random_crop_size is not None or y_random_crop_size is not None:
random_index = tf.random.uniform([],
minval=0,
maxval=num_points,
dtype=tf.int32)
center_x = points[random_index, 0]
center_y = points[random_index, 1]
points_x = points[:, 0]
points_y = points[:, 1]
min_x = tf.reduce_min(points_x) - epsilon
max_x = tf.reduce_max(points_x) + epsilon
min_y = tf.reduce_min(points_y) - epsilon
max_y = tf.reduce_max(points_y) + epsilon
if x_random_crop_size is not None:
min_x = center_x - x_random_crop_size / 2.0 - epsilon
max_x = center_x + x_random_crop_size / 2.0 + epsilon
if y_random_crop_size is not None:
min_y = center_y - y_random_crop_size / 2.0 - epsilon
max_y = center_y + y_random_crop_size / 2.0 + epsilon
x_mask = tf.logical_and(tf.greater(points_x, min_x),
tf.less(points_x, max_x))
y_mask = tf.logical_and(tf.greater(points_y, min_y),
tf.less(points_y, max_y))
points_mask = tf.logical_and(x_mask, y_mask)
for key in sorted(mesh_inputs):
mesh_inputs[key] = tf.boolean_mask(mesh_inputs[key], points_mask)
for key in sorted(view_indices_2d_inputs):
view_indices_2d_inputs[key] = tf.transpose(
tf.boolean_mask(
tf.transpose(view_indices_2d_inputs[key], [1, 0, 2]),
points_mask), [1, 0, 2])
def fit_gaussian_mixture(embeddings,
responsibilities,
damping=1e-7,
full_covariance=False):
"""Fits a unimodal Gaussian distribution `embeddings`.
Args:
embeddings: A [batch_size, embedding_dim] tf.Tensor of embeddings.
responsibilities: The per-component responsibilities.
damping: The scale of the covariance damping coefficient.
full_covariance: Whether to use a full or diagonal covariance.
Returns:
Parameter estimates for a Gaussian mixture model.
"""
num, dim = tf.split(tf.shape(input=embeddings), num_or_size_splits=2)
num, dim = tf.squeeze(num), tf.squeeze(dim)
num_classes = responsibilities.shape[1]
mixing_proportion = tf.einsum('jk->k', responsibilities)
mixing_proportion /= tf.cast(num, dtype=tf.float32)
mixing_logits = tf.math.log(mixing_proportion)
sample_mean = tf.einsum('ij,ik->jk', responsibilities, embeddings)
sample_mean /= tf.reduce_sum(
input_tensor=responsibilities, axis=0)[:, tf.newaxis]
centered_embeddings = (
embeddings[:, tf.newaxis, :] - sample_mean[tf.newaxis, :, :])
if full_covariance:
sample_covariance = tf.einsum('ijk,ijl->ijkl', centered_embeddings,
centered_embeddings) # Outer product.
sample_covariance += damping * tf.eye(dim) # Positive definiteness.
weighted_covariance = tf.einsum('ij,ijkl->jkl', responsibilities,
sample_covariance)
weighted_covariance /= tf.reduce_sum(
input_tensor=responsibilities, axis=0)[:, tf.newaxis, tf.newaxis]
return (
_split_and_squeeze(sample_mean, num_splits=num_classes),
_split_and_squeeze(weighted_covariance, num_splits=num_classes),
[mixing_logits],
)
else:
avg_x_squared = (
tf.matmul(responsibilities, embeddings**2, transpose_a=True) /
tf.reduce_sum(input_tensor=responsibilities, axis=0)[:, tf.newaxis])
avg_means_squared = sample_mean**2
avg_x_means = (
sample_mean *
tf.matmul(responsibilities, embeddings, transpose_a=True) /
tf.reduce_sum(input_tensor=responsibilities, axis=0)[:, tf.newaxis])
sample_variances = (
avg_x_squared - 2 * avg_x_means + avg_means_squared +
damping * tf.ones(dim))
log_variances = tf.math.log(sample_variances)
return (
_split_and_squeeze(sample_mean, num_splits=num_classes),
_split_and_squeeze(log_variances, num_splits=num_classes),
[mixing_logits],
)
def geometric_augmentation(images,
flow = None,
mask = None,
crop_height = 640,
crop_width = 640,
probability_flip_left_right = 0.5,
probability_flip_up_down = 0.1,
probability_scale = 0.8,
probability_relative_scale = 0.,
probability_stretch = 0.8,
probability_rotation = 0.0,
probability_relative_rotation = 0.0,
probability_crop_offset = 0.0,
min_bound_scale = -0.2,
max_bound_scale = 0.6,
max_strech_scale = 0.2,
min_bound_relative_scale = -0.1,
max_bound_relative_scale = 0.1,
max_rotation_deg = 15,
max_relative_rotation_deg = 3,
max_relative_crop_offset = 5,
return_full_scale=False):
"""Applies geometric augmentations to an image pair and corresponding flow.
Args:
images: Image pair of shape [2, height, width, channels].
flow: Corresponding forward flow field of shape [height, width, 2].
mask: Mask indicating which positions in the flow field hold valid flow
vectors of shape [height, width, 1]. Non-valid poisitions are encoded with
0, valid positions with 1.
crop_height: Height of the final augmented output.
crop_width: Width of the final augmented output.
probability_flip_left_right: Probability of applying left/right flip.
probability_flip_up_down: Probability of applying up/down flip
probability_scale: Probability of applying scale augmentation.
probability_relative_scale: Probability of applying scale augmentation to
only the second frame of the the image pair.
probability_stretch: Probability of applying stretch augmentation (scale
without keeping the aspect ratio).
probability_rotation: Probability of applying rotation augmentation.
probability_relative_rotation: Probability of applying rotation augmentation
to only the second frame of the the image pair.
probability_crop_offset: Probability of applying a relative offset while
cropping.
min_bound_scale: Defines the smallest possible scaling factor as
2**min_bound_scale.
max_bound_scale: Defines the largest possible scaling factor as
2**max_bound_scale.
max_strech_scale: Defines the smallest and largest possible streching factor
as 2**-max_strech_scale and 2**max_strech_scale.
min_bound_relative_scale: Defines the smallest possible scaling factor for
the relative scaling as 2**min_bound_relative_scale.
max_bound_relative_scale: Defines the largest possible scaling factor for
the relative scaling as 2**max_bound_relative_scale.
max_rotation_deg: Defines the maximum angle of rotation in degrees.
max_relative_rotation_deg: Defines the maximum angle of rotation in degrees
for the relative rotation.
max_relative_crop_offset: Defines the maximum relative offset in pixels for
cropping.
return_full_scale: bool. If this is passed, the full size images will be
returned in addition to the geometrically augmented (cropped and / or
resized) images. In addition to the resized images, the crop height,
width, and any padding applied will be returned.
Returns:
if return_full_scale is False:
Augmented images, flow and mask (if not None).
if return_full_scale is True:
Augmented images, flow, mask, full_size_images, crop_h, crop_w, pad_h,
and pad_w.
"""
# apply geometric augmentation
if probability_flip_left_right > 0:
images, flow, mask = random_flip_left_right(
images, flow, mask, probability_flip_left_right)
if probability_flip_up_down > 0:
images, flow, mask = random_flip_up_down(
images, flow, mask, probability_flip_up_down)
if probability_scale > 0 or probability_stretch > 0:
images, flow, mask = random_scale(
images,
flow,
mask,
min_scale=min_bound_scale,
max_scale=max_bound_scale,
max_strech=max_strech_scale,
probability_scale=probability_scale,
probability_strech=probability_stretch)
if probability_relative_scale > 0:
images, flow, mask = random_scale_second(
images, flow, mask,
min_scale=min_bound_relative_scale,
max_scale=max_bound_relative_scale,
probability_scale=probability_relative_scale)
if probability_rotation > 0:
images, flow, mask = random_rotation(
images, flow, mask,
probability=probability_rotation,
max_rotation=max_rotation_deg, not_empty_crop=True)
if probability_relative_rotation > 0:
images, flow, mask = random_rotation_second(
images, flow, mask,
probability=probability_relative_rotation,
max_rotation=max_relative_rotation_deg, not_empty_crop=True)
images_uncropped = images
images, flow, mask, offset_h, offset_w = random_crop(
images, flow, mask, crop_height, crop_width,
relative_offset=max_relative_crop_offset,
probability_crop_offset=probability_crop_offset)
# Add 100 / 200 pixels to crop height / width for full scale warp
pad_to_size_h = crop_height + 200
pad_to_size_w = crop_width + 400
if return_full_scale:
if pad_to_size_w:
uncropped_shape = tf.shape(images_uncropped)
if images.shape[1] > uncropped_shape[1] or images.shape[
2] > uncropped_shape[2]:
images_uncropped = images
uncropped_shape = tf.shape(images_uncropped)
offset_h = tf.zeros_like(offset_h)
offset_w = tf.zeros_like(offset_w)
if uncropped_shape[1] > pad_to_size_h:
crop_ht = offset_h - (200 // 2)
crop_hb = offset_h + crop_height + (200 // 2)
crop_hb += tf.maximum(0, -crop_ht)
crop_ht -= tf.maximum(0, -(uncropped_shape[1] - crop_hb))
crop_ht = tf.maximum(crop_ht, 0)
crop_hb = tf.minimum(crop_hb, uncropped_shape[1])
offset_h -= crop_ht
images_uncropped = images_uncropped[:, crop_ht:crop_hb, :, :]
if uncropped_shape[2] > pad_to_size_w:
crop_wt = offset_w - (400 // 2)
crop_wb = offset_w + crop_width + (400 // 2)
crop_wb += tf.maximum(0, -crop_wt)
crop_wt -= tf.maximum(0, -(uncropped_shape[2] - crop_wb))
crop_wt = tf.maximum(crop_wt, 0)
crop_wb = tf.minimum(crop_wb, uncropped_shape[2])
offset_w -= crop_wt
images_uncropped = images_uncropped[:, :, crop_wt:crop_wb, :]
uncropped_shape = tf.shape(images_uncropped)
# remove remove_pixels_w from the width while keeping the crop centered
pad_h = pad_to_size_h - uncropped_shape[1]
pad_w = pad_to_size_w - uncropped_shape[2]
with tf.control_dependencies([
tf.compat.v1.assert_greater_equal(pad_h, 0),
tf.compat.v1.assert_greater_equal(pad_w, 0)
]):
images_uncropped = tf.pad(images_uncropped,
[[0, 0], [pad_h, 0], [pad_w, 0], [0, 0]])
images_uncropped = tf.ensure_shape(images_uncropped,
[2, pad_to_size_h, pad_to_size_w, 3])
return images, flow, mask, images_uncropped, offset_h, offset_w, pad_h, pad_w
return images, flow, mask
def build_graph(self):
"""Builds the neural network graph."""
# define graph
self.g = tf.Graph()
with self.g.as_default():
# create and store a new session for the graph
self.sess = tf.Session()
# define placeholders
self.x = tf.placeholder(shape=[None, self.dim_input],
dtype=tf.float32)
self.y = tf.placeholder(shape=[None, self.num_classes],
dtype=tf.float32)
# linear layer(WX + b)
with tf.variable_scope('last_layer/dense') as scope:
weights = tf.get_variable('kernel',
[self.dim_input, self.num_classes],
dtype=tf.float32)
biases = tf.get_variable('bias', [self.num_classes],
dtype=tf.float32)
wb = tf.concat([weights, tf.expand_dims(biases, axis=0)], 0)
wb_renorm = tf.matmul(self.sigma_half_inv, wb)
weights_renorm = wb_renorm[:self.dim_input, :]
biases_renorm = wb_renorm[-1, :]
self.z = tf.add(tf.matmul(self.x, weights_renorm),
biases_renorm,
name=scope.name)
# Gaussian prior
# prior = tf.nn.l2_loss(weights) + tf.nn.l2_loss(biases)
# Non normalized loss, because of the preconditioning
self.loss = self.n * tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits_v2(labels=self.y,
logits=self.z))
# Bayesian loss
self.bayesian_loss = self.loss # + prior
self.output_probs = tf.nn.softmax(self.z)
# Variables of the last layer
self.ll_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
self.ll_vars_concat = tf.concat(
[self.ll_vars[0],
tf.expand_dims(self.ll_vars[1], axis=0)], 0)
# Summary
_variable_summaries(self.ll_vars_concat)
# saving the weights of last layer when running SGLD/SGD/MCMC algorithm
self.saver = tf.train.Saver(var_list=self.ll_vars,
max_to_keep=self.num_samples)
self.gd_opt = tf.train.GradientDescentOptimizer(self.step_size)
# SGLD optimizer for the last layer
if self.sampler in ['sgld', 'lmc']:
grads_vars = self.gd_opt.compute_gradients(self.bayesian_loss)
grads_vars_sgld = []
for g, v in grads_vars:
if g is not None:
s = list(v.name)
s[v.name.rindex(':')] = '_'
# Adding Gaussian noise to the gradient
gaussian_noise = (np.sqrt(2. / self.step_size) *
tf.random_normal(tf.shape(g)))
g_sgld = g + gaussian_noise
tf.summary.histogram(''.join(s) + '/grad_hist_mcmc', g)
tf.summary.histogram(
''.join(s) + '/gaussian_noise_hist_mcmc',
gaussian_noise)
tf.summary.histogram(
''.join(s) + '/grad_total_hist_mcmc', g_sgld)
grads_vars_sgld.append((g_sgld, v))
self.train_op = self.gd_opt.apply_gradients(grads_vars_sgld)
# SGD optimizer for the last layer
if self.sampler == 'sgd':
grads_vars_sgd = self.gd_opt.compute_gradients(self.loss)
self.train_op = self.gd_opt.apply_gradients(grads_vars_sgd)
for g, v in grads_vars_sgd:
if g is not None:
s = list(v.name)
s[v.name.rindex(':')] = '_'
tf.summary.histogram(''.join(s) + '/grad_hist_sgd', g)
# Merge all the summaries and write them out
self.all_summaries = tf.summary.merge_all()
location = os.path.join(self.working_dir, 'logs')
self.writer = tf.summary.FileWriter(location, graph=self.g)
saver_network = tf.train.Saver(var_list=self.ll_vars)
print('loading the network ...')
# Restores from checkpoint
saver_network.restore(self.sess, self.model_dir)
print('Graph successfully loaded.')
def true_fn(images, flow, mask):
angle_radian = tf.random.uniform(
[], minval=-max_rotation, maxval=max_rotation,
dtype=tf.float32) * math.pi / 180.0
image_1, image_2 = tf.unstack(images)
image_2 = rotate(image_2, angle_radian, is_flow=False, mask=None)
images = tf.stack([image_1, image_2])
if not_empty_crop:
orig_height = tf.shape(images)[-3]
orig_width = tf.shape(images)[-2]
# introduce abbreviations for shorter notation
cos = tf.math.cos(angle_radian % math.pi)
sin = tf.math.sin(angle_radian % math.pi)
h = tf.cast(orig_height, tf.float32)
w = tf.cast(orig_width, tf.float32)
# compute required scale factor
scale = tf.cond(tf.math.less(angle_radian % math.pi, math.pi/2.0),
lambda: tf.math.maximum((w/h)*sin+cos, (h/w)*sin+cos),
lambda: tf.math.maximum((w/h)*sin-cos, (h/w)*sin-cos))
new_height = tf.math.floor(h / scale)
new_width = tf.math.floor(w / scale)
# crop image again to original size
offset_height = tf.cast((h-new_height)/2, tf.int32)
offset_width = tf.cast((w-new_width)/2, tf.int32)
images = tf.image.crop_to_bounding_box(
images,
offset_height=offset_height,
offset_width=offset_width,
target_height=tf.cast(new_height, tf.int32),
target_width=tf.cast(new_width, tf.int32))
if flow is not None:
# get current locations (with the origin in the image center)
positions = _positions_center_origin(orig_height, orig_width)
# compute augmented flow (multiply by mask to zero invalid flow locations)
cos = tf.math.cos(angle_radian)
sin = tf.math.sin(angle_radian)
rotation_matrix = tf.reshape([cos, sin, -sin, cos], [2, 2])
flow = (tf.linalg.matmul(
(positions + flow), rotation_matrix) - positions) * mask
if not_empty_crop:
# crop flow and mask again to original size
flow = tf.image.crop_to_bounding_box(
flow,
offset_height=offset_height,
offset_width=offset_width,
target_height=tf.cast(new_height, tf.int32),
target_width=tf.cast(new_width, tf.int32))
mask = tf.image.crop_to_bounding_box(
mask,
offset_height=offset_height,
offset_width=offset_width,
target_height=tf.cast(new_height, tf.int32),
target_width=tf.cast(new_width, tf.int32))
return images, flow, mask
def prepare_lidar_images_and_correspondences(
inputs,
resized_image_height,
resized_image_width,
camera_names=('front', 'front_left', 'front_right', 'side_left',
'side_right'),
lidar_names=('top', 'front', 'side_left', 'side_right', 'rear')):
"""Integrates and returns the lidars, cameras and their correspondences.
Args:
inputs: A dictionary containing the images and point / pixel
correspondences.
resized_image_height: Target height of the images.
resized_image_width: Target width of the images.
camera_names: List of cameras to include images from.
lidar_names: List of lidars to include point clouds from.
Returns:
A tf.float32 tensor of size [num_points, 3] containing point positions.
A tf.float32 tensor of size [num_points, 1] containing point intensities.
A tf.float32 tensor of size [num_points, 1] containing point elongations.
A tf.float32 tensor of size [num_points, 3] containing point normals.
A tf.float32 tensor of size [num_images, resized_image_height,
resized_image_width, 3].
A tf.int32 tensor of size [num_images, num_points, 2].
Raises:
ValueError: If camera_names or lidar_names are empty lists.
"""
if not camera_names:
raise ValueError('camera_names should contain at least one name.')
if not lidar_names:
raise ValueError('lidar_names should contain at least one name.')
(points_position, points_intensity, points_elongation, points_normal,
points_in_image_frame_yx, points_in_image_frame_id) = _prepare_lidar_points(
inputs=inputs, lidar_names=lidar_names)
images = []
points_in_image_frame = []
for camera_name in camera_names:
image_key = ('cameras/%s/image' % camera_name)
image_height = tf.shape(inputs[image_key])[0]
image_width = tf.shape(inputs[image_key])[1]
height_ratio = tf.cast(
resized_image_height, dtype=tf.float32) / tf.cast(
image_height, dtype=tf.float32)
width_ratio = tf.cast(
resized_image_width, dtype=tf.float32) / tf.cast(
image_width, dtype=tf.float32)
if tf.executing_eagerly():
resize_method = tf.image.ResizeMethod.NEAREST_NEIGHBOR
else:
resize_method = tf.image.ResizeMethod.BILINEAR
if inputs[image_key].dtype in [
tf.int8, tf.uint8, tf.int16, tf.uint16, tf.int32, tf.int64
]:
resize_method = tf.image.ResizeMethod.NEAREST_NEIGHBOR
images.append(
tf.image.resize(
images=inputs[image_key],
size=[resized_image_height, resized_image_width],
method=resize_method,
antialias=True))
camera_id = tf.cast(inputs[('cameras/%s/id' % camera_name)], dtype=tf.int32)
valid_points = tf.equal(points_in_image_frame_id, camera_id)
valid_points = tf.tile(valid_points, [1, 2])
point_coords = tf.cast(
tf.cast(points_in_image_frame_yx, dtype=tf.float32) *
tf.stack([height_ratio, width_ratio]),
dtype=tf.int32)
points_in_image_frame_camera = tf.where(
valid_points, point_coords, -tf.ones_like(valid_points, dtype=tf.int32))
points_in_image_frame.append(points_in_image_frame_camera)
num_images = len(images)
images = tf.stack(images, axis=0)
images.set_shape([num_images, resized_image_height, resized_image_width, 3])
points_in_image_frame = tf.stack(points_in_image_frame, axis=0)
return {
'points_position': points_position,
'points_intensity': points_intensity,
'points_elongation': points_elongation,
'points_normal': points_normal,
'view_images': {'rgb_view': images},
'view_indices_2d': {'rgb_view': points_in_image_frame}
}
def points_to_normals_unbatched(points,
k,
distance_upper_bound,
viewpoint=None,
noise_magnitude=1e-4,
method='pca'):
"""Computes normals for the points in a point cloud.
Args:
points: A tf.float32 tensor of size [N, 3].
k: An integer determining the size of the neighborhood.
distance_upper_bound: Maximum distance of the neighbor points. If None, it
will not add a cap on the distance.
viewpoint: A tf.float32 tensor of size [3]. Normals will be flipped to point
towards view point. If None, it won't be used.
noise_magnitude: Noise magnitude to be added to the input of svd. If None,
it won't add noise.
method: The normal prediction method, options are `pca` and `cross` (cross
product).
Returns:
normals: A tf.float32 tensor of size [N, 3].
"""
if method == 'pca':
if k <= 3:
raise ValueError(
'At least 3 neighbors are required for computing PCA.')
elif method == 'cross':
if k <= 2:
raise ValueError(
'At least 2 neighbors are required for computing cross.')
else:
raise ValueError(('Unknown method of normal prediction %s' % method))
n = tf.shape(points)[0]
d = points.get_shape().as_list()[1]
if d != 3:
raise ValueError('Points dimension is not 3.')
_, knn_adjacencies = knn_graph_from_points_unbatched(
points=points, k=k, distance_upper_bound=distance_upper_bound)
knn_adjacencies = knn_adjacencies[:, 1:]
knn_adjacencies = tf.reshape(knn_adjacencies, [n * (k - 1)])
adjacency_points = tf.gather(points, indices=knn_adjacencies)
adjacency_points = tf.reshape(adjacency_points, [n, (k - 1), d])
if method == 'pca':
adjacency_relative_points = adjacency_points - tf.expand_dims(points,
axis=1)
if noise_magnitude is not None:
adjacency_relative_points += tf.random.uniform(
tf.shape(adjacency_relative_points),
minval=-noise_magnitude,
maxval=noise_magnitude,
dtype=tf.float32)
_, _, v = tf.linalg.svd(adjacency_relative_points)
normals = v[:, 2, :]
elif method == 'cross':
v1 = adjacency_points[:, 0, :] - points
v2 = adjacency_points[:, 1, :] - points
normals = tf.linalg.cross(v1, v2)
normals_length = tf.expand_dims(tf.norm(normals, axis=1), axis=1)
if noise_magnitude is not None:
normals_length += noise_magnitude
normals /= normals_length
else:
raise ValueError(('Unknown method of normal prediction %s' % method))
if viewpoint is not None:
normals = flip_normals_towards_viewpoint(points=points,
normals=normals,
viewpoint=viewpoint)
return normals
def build_graph(self):
"""Builds the neural network graph."""
# define graph
self.g = tf.Graph()
with self.g.as_default():
# create and store a new session for the graph
self.sess = tf.Session()
# define placeholders
self.x = tf.placeholder(shape=[None, self.dim_input],
dtype=tf.float32)
self.y = tf.placeholder(shape=[None, self.num_classes],
dtype=tf.float32)
# define simple model
with tf.variable_scope('last_layer'):
self.z = tf.layers.dense(inputs=self.x, units=self.num_classes)
self.loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits_v2(labels=self.y,
logits=self.z))
self.output_probs = tf.nn.softmax(self.z)
# Variables of the last layer
self.ll_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
self.ll_vars_concat = tf.concat(
[self.ll_vars[0],
tf.expand_dims(self.ll_vars[1], axis=0)], 0)
# Summary
_variable_summaries(self.ll_vars_concat)
# add regularization that acts as a unit Gaussian prior on the last layer
regularizer = tf.contrib.layers.l2_regularizer(1.0)
# regularization
prior = tf.contrib.layers.apply_regularization(
regularizer, self.ll_vars)
self.bayesian_loss = self.n * self.loss + prior
# saving the weights of last layer when running SGLD/SGD/MCMC algorithm
self.saver = tf.train.Saver(var_list=self.ll_vars,
max_to_keep=self.num_samples)
# SGLD optimizer for the last layer
if self.sampler in ['sgld', 'lmc']:
step = self.step_size / self.n
gd_opt = tf.train.GradientDescentOptimizer(step)
grads_vars = gd_opt.compute_gradients(self.bayesian_loss)
grads_vars_sgld = []
for g, v in grads_vars:
if g is not None:
s = list(v.name)
s[v.name.rindex(':')] = '_'
# Adding Gaussian noise to the gradient
gaussian_noise = (np.sqrt(2. / step) *
tf.random_normal(tf.shape(g)))
g_sgld = g + gaussian_noise
tf.summary.histogram(''.join(s) + '/grad_hist_mcmc',
g / self.n)
tf.summary.histogram(
''.join(s) + '/gaussian_noise_hist_mcmc',
gaussian_noise / self.n)
tf.summary.histogram(
''.join(s) + '/grad_total_hist_mcmc',
g_sgld / self.n)
grads_vars_sgld.append((g_sgld, v))
self.train_op = gd_opt.apply_gradients(grads_vars_sgld)
# SGD optimizer for the last layer
if self.sampler == 'sgd':
gd_opt = tf.train.GradientDescentOptimizer(self.step_size)
grads_vars_sgd = gd_opt.compute_gradients(self.loss)
self.train_op = gd_opt.apply_gradients(grads_vars_sgd)
for g, v in grads_vars_sgd:
if g is not None:
s = list(v.name)
s[v.name.rindex(':')] = '_'
tf.summary.histogram(''.join(s) + '/grad_hist_sgd', g)
# Merge all the summaries and write them out
self.all_summaries = tf.summary.merge_all()
location = os.path.join(self.working_dir, 'logs')
self.writer = tf.summary.FileWriter(location, graph=self.g)
saver_network = tf.train.Saver(var_list=self.ll_vars)
print('loading the network ...')
# Restores from checkpoint
# self.sess.run(tf.global_variables_initializer())
saver_network.restore(self.sess, self.model_dir)
print('Graph successfully loaded.')
