def flip_randomly_points_and_normals_motions(points, normals, motions,
is_training):
"""Flip points and normals against x or/and y axis.
Args:
points: A tf.float32 tensor of size [N, 3] containing points.
normals: A tf.float32 tensor of size [N, 3] containing points or None.
motions: A tf.float32 tensor of size [N, 3] containing motion vectors or
None.
is_training: True if in training stage. Random flipping only takes place
during training.
Returns:
flipped_points: Flipped points. A tf.float32 tensor of size [N, 3].
flipped_normals: Flipped normals. A tf.float32 tensor of size [N, 3]. It
will be None of the normals is None.
"""
if is_training:
x_cond = tf.greater(
tf.random.uniform([], minval=0.0, maxval=1.0, dtype=tf.float32),
0.5)
x_rotate = tf.cond(x_cond, lambda: tf.constant(1.0, dtype=tf.float32),
lambda: tf.constant(-1.0, dtype=tf.float32))
y_cond = tf.greater(
tf.random.uniform([], minval=0.0, maxval=1.0, dtype=tf.float32),
0.5)
y_rotate = tf.cond(y_cond, lambda: tf.constant(1.0, dtype=tf.float32),
lambda: tf.constant(-1.0, dtype=tf.float32))
(points, normals,
motions) = flip_points_and_normals_motions(points=points,
normals=normals,
motions=motions,
x_rotate=x_rotate,
y_rotate=y_rotate)
return points, normals, motions
def _get_random_scaled_resolution(orig_height, orig_width, min_scale,
max_scale, max_strech, probability_strech):
"""Computes a new random resolution."""
# Choose a random scale factor and compute new resolution.
scale = 2**tf.random.uniform([],
minval=min_scale,
maxval=max_scale,
dtype=tf.float32)
scale_height = scale
scale_width = scale
# Possibly change scale values individually to perform strech
def true_fn(scale_height, scale_width):
scale_height *= 2**tf.random.uniform([], -max_strech, max_strech)
scale_width *= 2**tf.random.uniform([], -max_strech, max_strech)
return tf.stack((scale_height, scale_width), axis=0)
def false_fn(scale_height, scale_width):
return tf.stack((scale_height, scale_width), axis=0)
perform_strech = tf.random.uniform([]) < probability_strech
scales = tf.cond(perform_strech,
lambda: true_fn(scale_height, scale_width),
lambda: false_fn(scale_height, scale_width))
scale_height = scales[0]
scale_width = scales[1]
# Compute scaled image resolution.
new_height = tf.cast(
tf.math.ceil(tf.cast(orig_height, tf.float32) * scale_height),
tf.int32)
new_width = tf.cast(
tf.math.ceil(tf.cast(orig_width, tf.float32) * scale_width), tf.int32)
return new_height, new_width, scale
def _set_up_staging(self, transition):
"""Sets up staging ops for prefetching the next transition.
This allows us to hide the py_func latency. To do so we use a staging area
to pre-fetch the next batch of transitions.
Args:
transition: tuple of tf.Tensors with shape
memory.get_transition_elements().
Returns:
prefetched_transition: tuple of tf.Tensors with shape
memory.get_transition_elements() that have been previously prefetched.
"""
transition_type = self.memory.get_transition_elements()
# Create the staging area in CPU.
prefetch_area = tf.contrib.staging.StagingArea(
[shape_with_type.type for shape_with_type in transition_type])
# Store prefetch op for tests, but keep it private -- users should not be
# calling _prefetch_batch.
self._prefetch_batch = prefetch_area.put(transition)
initial_prefetch = tf.cond(
tf.equal(prefetch_area.size(), 0),
lambda: prefetch_area.put(transition), tf.no_op)
# Every time a transition is sampled self.prefetch_batch will be
# called. If the staging area is empty, two put ops will be called.
with tf.control_dependencies([self._prefetch_batch, initial_prefetch]):
prefetched_transition = prefetch_area.get()
return prefetched_transition
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
def _box_center_distance_loss_on_voxel_tensors_unbatched(
inputs_1, outputs_1, loss_type, delta, is_balanced, is_intermediate):
"""Computes huber loss on predicted object centers for each voxel."""
inputs_1, outputs_1, valid_mask = _get_voxels_valid_inputs_outputs(
inputs_1=inputs_1, outputs_1=outputs_1)
def loss_fn_unbatched():
"""Loss function."""
if is_intermediate:
output_boxes_center = outputs_1[standard_fields.DetectionResultFields
.intermediate_object_center_voxels]
else:
output_boxes_center = outputs_1[
standard_fields.DetectionResultFields.object_center_voxels]
return _box_center_distance_loss(
loss_type=loss_type,
is_balanced=is_balanced,
input_boxes_center=inputs_1[
standard_fields.InputDataFields.object_center_voxels],
input_boxes_instance_id=inputs_1[
standard_fields.InputDataFields.object_instance_id_voxels],
output_boxes_center=output_boxes_center,
delta=delta)
return tf.cond(
tf.reduce_any(valid_mask),
loss_fn_unbatched, lambda: tf.constant(0.0, dtype=tf.float32))
def random_scale(images,
flow,
mask,
min_scale,
max_scale,
max_strech,
probability_scale,
probability_strech):
"""Performs a random scaling in the given range."""
perform_scale = tf.random.uniform([]) < probability_scale
def true_fn(images, flow, mask):
# Get a random new resolution to which the images will be scaled.
orig_height = tf.shape(images)[-3]
orig_width = tf.shape(images)[-2]
new_height, new_width, _ = _get_random_scaled_resolution(
orig_height=orig_height,
orig_width=orig_width,
min_scale=min_scale,
max_scale=max_scale,
max_strech=max_strech,
probability_strech=probability_strech)
# rescale the images (and flow)
images = smurf_utils.resize(images, new_height, new_width, is_flow=False)
if flow is not None:
flow, mask = smurf_utils.resize(
flow, new_height, new_width, is_flow=True, mask=mask)
return images, flow, mask
def false_fn(images, flow, mask):
return images, flow, mask
return tf.cond(perform_scale, lambda: true_fn(images, flow, mask),
lambda: false_fn(images, flow, mask))
def _train_op(self):
with tf.name_scope("train_op"):
d_opt = tf.train.GradientDescentOptimizer(self.d_lr)
var_list = tf.trainable_variables(self.scope + "/discriminator")
gvs, d_norm = clip_grads(self.d_loss, var_list)
self.d_train = d_opt.minimize(self.d_loss,
var_list=var_list,
global_step=self._global_step)
g_opt = tf.train.AdamOptimizer(self.g_lr)
var_list = tf.trainable_variables(self.scope + "/generator")
gvs, g_norm = clip_grads(self.g_loss, var_list)
self.g_train = g_opt.minimize(self.g_loss,
var_list=var_list,
global_step=self._global_step)
# g_train = g_opt.apply_gradients(gvs, global_step=self._global_step)
self.train = tf.cond(self.flag, lambda: self.g_train,
lambda: self.d_train)
self._summary_dict.update({
"distance":
self._gen_norm(self.x_fake, self.y),
"g_norm":
g_norm,
"d_norm":
d_norm,
"g_loss":
self.g_loss,
"d_loss":
self.d_loss
})
def _box_rotation_regression_loss(loss_type, is_balanced,
input_boxes_rotation_matrix,
input_boxes_instance_id,
output_boxes_rotation_matrix, delta):
"""Computes regression loss on object rotations."""
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]))
cond_input = tf.greater(tf.shape(input_boxes_rotation_matrix)[0], 0)
cond_output = tf.greater(tf.shape(output_boxes_rotation_matrix)[0], 0)
cond = tf.logical_and(cond_input, cond_output)
return tf.cond(cond, fn, lambda: tf.constant(0.0, dtype=tf.float32))
def pick_labeled_image(mesh_inputs, view_image_inputs, view_indices_2d_inputs,
view_name):
"""Pick the image with most number of labeled points projecting to it."""
if view_name not in view_image_inputs:
return
if view_name not in view_indices_2d_inputs:
return
if standard_fields.InputDataFields.point_loss_weights not in mesh_inputs:
raise ValueError('The key `weights` is missing from mesh_inputs.')
height = tf.shape(view_image_inputs[view_name])[1]
width = tf.shape(view_image_inputs[view_name])[2]
valid_points_y = tf.logical_and(
tf.greater_equal(view_indices_2d_inputs[view_name][:, :, 0], 0),
tf.less(view_indices_2d_inputs[view_name][:, :, 0], height))
valid_points_x = tf.logical_and(
tf.greater_equal(view_indices_2d_inputs[view_name][:, :, 1], 0),
tf.less(view_indices_2d_inputs[view_name][:, :, 1], width))
valid_points = tf.logical_and(valid_points_y, valid_points_x)
image_total_weights = tf.reduce_sum(
tf.cast(valid_points, dtype=tf.float32) * tf.squeeze(
mesh_inputs[standard_fields.InputDataFields.point_loss_weights],
axis=1),
axis=1)
image_total_weights = tf.cond(
tf.equal(tf.reduce_sum(image_total_weights), 0),
lambda: tf.reduce_sum(tf.cast(valid_points, dtype=tf.float32), axis=1),
lambda: image_total_weights)
best_image = tf.math.argmax(image_total_weights)
view_image_inputs[view_name] = view_image_inputs[view_name][
best_image:best_image + 1, :, :, :]
view_indices_2d_inputs[view_name] = view_indices_2d_inputs[view_name][
best_image:best_image + 1, :, :]
def random_brightness(images,
probability,
max_delta):
perform_brightness = tf.random.uniform([]) < probability
return tf.cond(
perform_brightness,
lambda: tf.image.random_brightness(images, max_delta),
lambda: images)
def random_rotation(images,
flow=None,
mask=None,
max_rotation=10,
not_empty_crop=True):
"""Performs a random rotation with the specified maximum rotation."""
angle_radian = tf.random.uniform(
[], minval=-max_rotation, maxval=max_rotation,
dtype=tf.float32) * pi / 180.0
images = rotate(images, angle_radian, is_flow=False, mask=None)
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 % pi)
sin = tf.math.sin(angle_radian % 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 % pi, 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:
flow, mask = rotate(flow, angle_radian, is_flow=True, mask=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 random_saturation(images,
probability,
min_bound,
max_bound):
perform_saturation = tf.random.uniform([]) < probability
return tf.cond(
perform_saturation,
lambda: tf.image.random_saturation(images, min_bound, max_bound),
lambda: images)
def random_contrast(images,
probability,
min_bound,
max_bound):
perform_contrast = tf.random.uniform([]) < probability
return tf.cond(
perform_contrast,
lambda: tf.image.random_contrast(images, min_bound, max_bound),
lambda: images)
def add_simclr_episodes(simclr_episode_fraction, *episode):
"""Convert simclr_episode_fraction of episodes into SimCLR Episodes."""
def convert_to_simclr_episode(support_images=None,
support_labels=None,
support_class_ids=None,
query_images=None,
query_labels=None,
query_class_ids=None):
"""Convert a single episode into a SimCLR Episode."""
# If there were k query examples of class c, keep the first k support
# examples of class c as 'simclr' queries. We do this by assigning an
# id for each image in the query set, implemented as label*1e5+x+1, where
# x is the number of images of the same label with a lower index within
# the query set. We do the same for the support set, which gives us a
# mapping between query and support images which is injective (as long
# as there's enough support-set images of each class).
#
# note: assumes max support label is 10000 - max_images_per_class
query_idx_within_class = tf.cast(
tf.equal(query_labels[tf.newaxis, :], query_labels[:, tf.newaxis]),
tf.int32)
query_idx_within_class = tf.linalg.diag_part(
tf.cumsum(query_idx_within_class, axis=1))
query_uid = query_labels * 10000 + query_idx_within_class
support_idx_within_class = tf.cast(
tf.equal(support_labels[tf.newaxis, :],
support_labels[:, tf.newaxis]), tf.int32)
support_idx_within_class = tf.linalg.diag_part(
tf.cumsum(support_idx_within_class, axis=1))
support_uid = support_labels * 10000 + support_idx_within_class
# compute which support-set images have matches in the query set, and
# discard the rest to produce the new query set.
support_keep = tf.reduce_any(tf.equal(support_uid[:, tf.newaxis],
query_uid[tf.newaxis, :]),
axis=1)
query_images = tf.boolean_mask(support_images, support_keep)
support_labels = tf.range(tf.shape(support_labels)[0],
dtype=support_labels.dtype)
query_labels = tf.boolean_mask(support_labels, support_keep)
query_class_ids = tf.boolean_mask(support_class_ids, support_keep)
# Finally, apply SimCLR augmentation to all images.
# Note simclr only blurs one image.
query_images = simclr_augment(query_images, blur=True)
support_images = simclr_augment(support_images)
return (support_images, support_labels, support_class_ids,
query_images, query_labels, query_class_ids)
return tf.cond(
tf.random_uniform([], minval=0, maxval=1,
dtype=tf.float32) > simclr_episode_fraction,
lambda: episode, lambda: convert_to_simclr_episode(*episode))
def random_scale_second(images,
flow,
mask,
min_scale,
max_scale,
probability_scale):
"""Performs a random scaling on the second image in the given range."""
perform_scale = tf.random.uniform([]) < probability_scale
def true_fn(images, flow, mask):
# choose a random scale factor and compute new resolution
orig_height = tf.shape(images)[-3]
orig_width = tf.shape(images)[-2]
new_height, new_width, scale = _get_random_scaled_resolution(
orig_height=orig_height,
orig_width=orig_width,
min_scale=min_scale,
max_scale=max_scale,
max_strech=0.0,
probability_strech=0.0)
# rescale only the second image
image_1, image_2 = tf.unstack(images)
image_2 = smurf_utils.resize(image_2, new_height, new_width, is_flow=False)
# Crop either first or second image to have matching dimensions
if scale < 1.0:
image_1 = _center_crop(image_1, new_height, new_width)
else:
image_2 = _center_crop(image_2, orig_height, orig_width)
images = tf.stack([image_1, image_2])
if flow is not None:
# get current locations (with the origin in the image center)
positions = _positions_center_origin(orig_height, orig_width)
# compute scale factor of the actual new image resolution
scale_flow_h = tf.cast(new_height, tf.float32) / tf.cast(
orig_height, tf.float32)
scale_flow_w = tf.cast(new_width, tf.float32) / tf.cast(
orig_width, tf.float32)
scale_flow = tf.stack([scale_flow_h, scale_flow_w])
# compute augmented flow (multiply by mask to zero invalid flow locations)
flow = ((positions + flow) * scale_flow - positions) * mask
if scale < 1.0:
# in case we downsample the image we crop the reference image to keep
# the same shape
flow = _center_crop(flow, new_height, new_width)
mask = _center_crop(mask, new_height, new_width)
return images, flow, mask
def false_fn(images, flow, mask):
return images, flow, mask
return tf.cond(perform_scale, lambda: true_fn(images, flow, mask),
lambda: false_fn(images, flow, mask))
def random_flip_left_right(images, flow, mask, probability):
"""Performs a random left/right flip."""
perform_flip = tf.less(tf.random.uniform([]), probability)
# apply flip
images = tf.cond(pred=perform_flip,
true_fn=lambda: tf.reverse(images, axis=[-2]),
false_fn=lambda: images)
if flow is not None:
flow = tf.cond(pred=perform_flip,
true_fn=lambda: tf.reverse(flow, axis=[-2]),
false_fn=lambda: flow)
mask = tf.cond(pred=perform_flip,
true_fn=lambda: tf.reverse(mask, axis=[-2]),
false_fn=lambda: mask)
# correct sign of flow
sign_correction = tf.reshape([1.0, -1.0], [1, 1, 2])
flow = tf.cond(pred=perform_flip,
true_fn=lambda: flow * sign_correction,
false_fn=lambda: flow)
return images, flow, mask
def random_flip_up_down(images, flow=None, mask=None):
"""Performs a random up/down flip."""
# 50/50 chance
perform_flip = tf.equal(tf.random.uniform([], maxval=2, dtype=tf.int32), 1)
# apply flip
images = tf.cond(pred=perform_flip,
true_fn=lambda: tf.reverse(images, axis=[-3]),
false_fn=lambda: images)
if flow is not None:
flow = tf.cond(pred=perform_flip,
true_fn=lambda: tf.reverse(flow, axis=[-3]),
false_fn=lambda: flow)
mask = tf.cond(pred=perform_flip,
true_fn=lambda: tf.reverse(mask, axis=[-3]),
false_fn=lambda: mask)
# correct sign of flow
sign_correction = tf.reshape([-1.0, 1.0], [1, 1, 2])
flow = tf.cond(pred=perform_flip,
true_fn=lambda: flow * sign_correction,
false_fn=lambda: flow)
return images, flow, mask
def reset_rewards(states,
actions,
rewards,
next_states,
contexts,
reset_index=0,
reset_state=None,
reset_reward_function=None,
include_forward_rewards=True,
include_reset_rewards=True):
"""Returns the rewards for a forward/reset agent.
Args:
states: A [batch_size, num_state_dims] Tensor representing a batch
of states.
actions: A [batch_size, num_action_dims] Tensor representing a batch
of actions.
rewards: A [batch_size] Tensor representing a batch of rewards.
next_states: A [batch_size, num_state_dims] Tensor representing a batch
of next states.
contexts: A list of [batch_size, num_context_dims] Tensor representing
a batch of contexts.
reset_index: (integer) The context list index that specifies reset.
reset_state: Reset state.
reset_reward_function: Reward function for reset step.
include_forward_rewards: Include the rewards from the forward pass.
include_reset_rewards: Include the rewards from the reset pass.
Returns:
A new tf.float32 [batch_size] rewards Tensor, and
tf.float32 [batch_size] discounts tensor.
"""
reset_state = tf.constant(
reset_state, dtype=next_states.dtype, shape=next_states.shape)
reset_states = tf.expand_dims(reset_state, 0)
def true_fn():
if include_reset_rewards:
return reset_reward_function(states, actions, rewards, next_states,
[reset_states] + contexts[1:])
else:
return tf.zeros_like(rewards), tf.ones_like(rewards)
def false_fn():
if include_forward_rewards:
return plain_rewards(states, actions, rewards, next_states, contexts)
else:
return tf.zeros_like(rewards), tf.ones_like(rewards)
rewards, discounts = tf.cond(
tf.cast(contexts[reset_index][0, 0], dtype=tf.bool), true_fn, false_fn)
return rewards, discounts
def random_color_swap(images, probability):
"""Randomly permute colors (rolling and reversing covers all permutations)."""
perform_color_swap = tf.random.uniform([]) < probability
def true_fn(images):
r = tf.random.uniform([], maxval=3, dtype=tf.int32)
images = tf.roll(images, r, axis=-1)
r = tf.equal(tf.random.uniform([], maxval=2, dtype=tf.int32), 1)
return tf.reverse(images, axis=[-1])
def false_fn(images):
return images
return tf.cond(perform_color_swap,
lambda: true_fn(images),
lambda: false_fn(images))
def random_contrast_individual(images, probability, min_bound, max_bound):
perform_augmentation = tf.random.uniform([]) < probability
def true_fn(images):
image_1, image_2 = tf.unstack(images)
image_1 = tf.image.random_contrast(image_1, min_bound, max_bound)
image_2 = tf.image.random_contrast(image_2, min_bound, max_bound)
return tf.stack([image_1, image_2])
def false_fn(images):
return images
return tf.cond(perform_augmentation, lambda: true_fn(images),
lambda: false_fn(images))
def random_brightness_individual(images,
probability,
max_delta):
perform_augmentation = tf.random.uniform([]) < probability
def true_fn(images):
image_1, image_2 = tf.unstack(images)
image_1 = tf.image.random_brightness(image_1, max_delta)
image_2 = tf.image.random_brightness(image_2, max_delta)
return tf.stack([image_1, image_2])
def false_fn(images):
return images
return tf.cond(perform_augmentation,
lambda: true_fn(images),
lambda: false_fn(images))
def box_corner_distance_loss_on_object_tensors(inputs,
outputs,
loss_type,
delta=1.0,
is_balanced=False):
"""Computes regression loss on object corner locations using object tensors.
Args:
inputs: A dictionary of tf.Tensors with our input data.
outputs: A dictionary of tf.Tensors with the network output.
loss_type: Loss type.
delta: float, the voxel where the huber loss function changes from a
quadratic to linear.
is_balanced: If True, the per-voxel losses are re-weighted to have equal
total weight for each object instance.
Returns:
localization_loss: A tf.float32 scalar corresponding to localization loss.
"""
def fn(inputs_1, outputs_1):
return _box_corner_distance_loss_on_object_tensors(
inputs=inputs_1,
outputs=outputs_1,
loss_type=loss_type,
delta=delta,
is_balanced=is_balanced)
batch_size = len(inputs[standard_fields.InputDataFields.objects_length])
losses = []
for b in range(batch_size):
inputs_1 = batch_utils.get_batch_size_1_input_objects(inputs=inputs,
b=b)
outputs_1 = batch_utils.get_batch_size_1_output_objects(
outputs=outputs, b=b)
cond_input = tf.greater(
tf.shape(
inputs_1[standard_fields.InputDataFields.objects_length])[0],
0)
cond_output = tf.greater(
tf.shape(outputs_1[
standard_fields.DetectionResultFields.objects_length])[0], 0)
cond = tf.logical_and(cond_input, cond_output)
# pylint: disable=cell-var-from-loop
loss = tf.cond(cond,
lambda: fn(inputs_1=inputs_1, outputs_1=outputs_1),
lambda: tf.constant(0.0, dtype=tf.float32))
# pylint: enable=cell-var-from-loop
losses.append(loss)
return tf.reduce_mean(tf.stack(losses))
def _box_corner_distance_loss_on_object_tensors(inputs, outputs, loss_type,
delta, is_balanced):
"""Computes huber loss on object corner locations."""
valid_mask_class = tf.greater(
tf.reshape(inputs[standard_fields.InputDataFields.objects_class],
[-1]), 0)
valid_mask_instance = tf.greater(
tf.reshape(inputs[standard_fields.InputDataFields.objects_instance_id],
[-1]), 0)
valid_mask = tf.logical_and(valid_mask_class, valid_mask_instance)
def fn():
for field in standard_fields.get_input_object_fields():
if field in inputs:
inputs[field] = tf.boolean_mask(inputs[field], valid_mask)
for field in standard_fields.get_output_object_fields():
if field in outputs:
outputs[field] = tf.boolean_mask(outputs[field], valid_mask)
return _box_corner_distance_loss(
loss_type=loss_type,
is_balanced=is_balanced,
input_boxes_length=inputs[
standard_fields.InputDataFields.objects_length],
input_boxes_height=inputs[
standard_fields.InputDataFields.objects_height],
input_boxes_width=inputs[
standard_fields.InputDataFields.objects_width],
input_boxes_center=inputs[
standard_fields.InputDataFields.objects_center],
input_boxes_rotation_matrix=inputs[
standard_fields.InputDataFields.objects_rotation_matrix],
input_boxes_instance_id=inputs[
standard_fields.InputDataFields.objects_instance_id],
output_boxes_length=outputs[
standard_fields.DetectionResultFields.objects_length],
output_boxes_height=outputs[
standard_fields.DetectionResultFields.objects_height],
output_boxes_width=outputs[
standard_fields.DetectionResultFields.objects_width],
output_boxes_center=outputs[
standard_fields.DetectionResultFields.objects_center],
output_boxes_rotation_matrix=outputs[
standard_fields.DetectionResultFields.objects_rotation_matrix],
delta=delta)
return tf.cond(tf.reduce_any(valid_mask), fn,
lambda: tf.constant(0.0, dtype=tf.float32))
def _build_train_op(self):
"""Builds a training op.
Returns:
An op performing one step of training from replay data.
"""
# click_indicator: [B, S]
# q_values: [B, A]
# actions: [B, S]
# slate_q_values: [B, S]
# replay_click_q: [B]
click_indicator = self._replay.rewards[:, :,
self._click_response_index]
slate_q_values = tf.compat.v1.batch_gather(
self._replay_net_outputs.q_values,
tf.cast(self._replay.actions, dtype=tf.int32))
# Only get the Q from the clicked document.
replay_click_q = tf.reduce_sum(input_tensor=slate_q_values *
click_indicator,
axis=1,
name='replay_click_q')
target = tf.stop_gradient(self._build_target_q_op())
clicked = tf.reduce_sum(input_tensor=click_indicator, axis=1)
clicked_indices = tf.squeeze(tf.compat.v1.where(tf.equal(clicked, 1)),
axis=1)
# clicked_indices is a vector and tf.gather selects the batch dimension.
q_clicked = tf.gather(replay_click_q, clicked_indices)
target_clicked = tf.gather(target, clicked_indices)
def get_train_op():
loss = tf.reduce_mean(input_tensor=tf.square(q_clicked -
target_clicked))
if self.summary_writer is not None:
with tf.compat.v1.variable_scope('Losses'):
tf.compat.v1.summary.scalar('Loss', loss)
return loss
loss = tf.cond(pred=tf.greater(tf.reduce_sum(input_tensor=clicked), 0),
true_fn=get_train_op,
false_fn=lambda: tf.constant(0.),
name='')
return self.optimizer.minimize(loss)
def random_gaussian_noise(images,
probability,
min_bound,
max_bound):
"""Augments images by adding gaussian noise."""
perform_gaussian_noise = tf.random.uniform([]) < probability
def true_fn(images):
sigma = tf.random.uniform([],
minval=min_bound,
maxval=max_bound,
dtype=tf.float32)
noise = tf.random.normal(
tf.shape(input=images), stddev=sigma, dtype=tf.float32)
images = images + noise
def false_fn(images):
return images
return tf.cond(perform_gaussian_noise,
lambda: true_fn(images),
lambda: false_fn(images))
def _box_size_regression_loss_on_voxel_tensors_unbatched(
inputs_1, outputs_1, loss_type, delta, is_balanced, is_intermediate):
"""Computes regression loss on predicted object size for each voxel."""
inputs_1, outputs_1, valid_mask = _get_voxels_valid_inputs_outputs(
inputs_1=inputs_1, outputs_1=outputs_1)
def loss_fn_unbatched():
"""Loss function."""
if is_intermediate:
output_boxes_length = outputs_1[
standard_fields.DetectionResultFields.
intermediate_object_length_voxels]
output_boxes_height = outputs_1[
standard_fields.DetectionResultFields.
intermediate_object_height_voxels]
output_boxes_width = outputs_1[
standard_fields.DetectionResultFields.
intermediate_object_width_voxels]
else:
output_boxes_length = outputs_1[
standard_fields.DetectionResultFields.object_length_voxels]
output_boxes_height = outputs_1[
standard_fields.DetectionResultFields.object_height_voxels]
output_boxes_width = outputs_1[
standard_fields.DetectionResultFields.object_width_voxels]
return _box_size_regression_loss(
loss_type=loss_type,
is_balanced=is_balanced,
input_boxes_length=inputs_1[
standard_fields.InputDataFields.object_length_voxels],
input_boxes_height=inputs_1[
standard_fields.InputDataFields.object_height_voxels],
input_boxes_width=inputs_1[
standard_fields.InputDataFields.object_width_voxels],
input_boxes_instance_id=inputs_1[
standard_fields.InputDataFields.object_instance_id_voxels],
output_boxes_length=output_boxes_length,
output_boxes_height=output_boxes_height,
output_boxes_width=output_boxes_width,
delta=delta)
return tf.cond(tf.reduce_any(valid_mask), loss_fn_unbatched,
lambda: tf.constant(0.0, dtype=tf.float32))
def _box_size_regression_loss(loss_type, is_balanced, input_boxes_length,
input_boxes_height, input_boxes_width,
input_boxes_instance_id, output_boxes_length,
output_boxes_height, output_boxes_width, delta):
"""Computes regression loss on object sizes."""
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_length = tf.reshape(input_boxes_length, [-1, 1])
gt_height = tf.reshape(input_boxes_height, [-1, 1])
gt_width = tf.reshape(input_boxes_width, [-1, 1])
predicted_length = tf.reshape(output_boxes_length, [-1, 1])
predicted_height = tf.reshape(output_boxes_height, [-1, 1])
predicted_width = tf.reshape(output_boxes_width, [-1, 1])
predicted_length /= gt_length
predicted_height /= gt_height
predicted_width /= gt_width
predicted_size = tf.concat(
[predicted_length, predicted_height, predicted_width], axis=1)
gt_size = tf.ones_like(predicted_size)
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))
size_losses = loss_fn(y_true=gt_size, y_pred=predicted_size)
return tf.reduce_mean(size_losses * tf.reshape(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 _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 inner_objective(self, onehot_labels, predictions, iteration_idx):
"""Compute the inner-loop objective."""
# p(z, y), joint log-likelihood.
joint_log_probs = self.joint_log_likelihood(onehot_labels, predictions)
labels = tf.expand_dims(tf.argmax(input=onehot_labels, axis=-1),
axis=-1)
numerator = tf.gather(joint_log_probs, labels, axis=-1, batch_dims=1)
# p(z), normalization constant.
evidence = tf.reduce_logsumexp(input_tensor=joint_log_probs,
axis=-1,
keepdims=True)
# p(y | z) if interpolation coefficient > 0 else p(z, y).
# TODO(eringrant): This assumes that `interp` is either 1 or 0.
# Adapt to a hybridized approach.
interp = tf.gather(self.gen_disc_interpolation, iteration_idx)
scale = tf.cond(pred=interp > 0.0,
true_fn=lambda: 1.0,
false_fn=lambda: self.generative_scaling)
return -scale * tf.reduce_mean(
input_tensor=numerator - interp * evidence, axis=0)
def random_hue_shift(images,
probability,
max_delta):
perform_hue_shift = tf.random.uniform([]) < probability
return tf.cond(perform_hue_shift,
lambda: tf.image.random_hue(images, max_delta), lambda: images)