def center_nn_loss(model_config, training_example, structured_implicit): # pylint:disable=unused-argument
"""A loss that decreases with the nearest neighbor center->center distance."""
# Training example present for interface uniformity
element_centers = structured_implicit.element_centers
center_shape = element_centers.get_shape().as_list()
if len(center_shape) != 3:
raise ValueError(f'Expected the element centers to have shape [b, #, 3],'
f' but they have shape {center_shape}. Loss=center_nn.')
batch_size, center_count, _ = center_shape
sq_distances = tf.reduce_sum(
tf.square(
tf.reshape(element_centers, [batch_size, center_count, 1, 3]) -
tf.reshape(element_centers, [batch_size, 1, center_count, 3])),
axis=-1)
distances = tf.sqrt(sq_distances + 1e-8)
loss_max = model_config.hparams.nnt
# We have to give the diagonal self -> self distances a high weight so they
# aren't valid choices:
diag_distances = tf.diag(tf.ones([center_count]) * (loss_max + 1))
diag_distances = tf.reshape(diag_distances, [1, center_count, center_count])
distances = distances + diag_distances
min_dists = tf.reduce_min(distances, axis=-1) # Shape [BS, #].
assert len(min_dists.shape) == 2
loss = tf.reduce_mean(tf.maximum(loss_max - min_dists,
0.0)) * model_config.hparams.nw
summarize.summarize_loss(model_config, loss, 'center-nn-loss')
return loss
def element_center_lowres_grid_direct_loss(model_config, training_example,
structured_implicit):
element_centers = structured_implicit.element_centers
gt_sdf_at_centers, _ = interpolate_util.interpolate(
training_example.grid, element_centers, training_example.world2grid)
mse = model_config.hparams.gd * tf.reduce_mean(gt_sdf_at_centers) + 1e-5
summarize.summarize_loss(model_config, mse, 'lowres_grid_direct_loss')
return mse
def element_center_lowres_grid_squared_loss(model_config, training_example,
structured_implicit):
element_centers = structured_implicit.element_centers
gt_sdf_at_centers, _ = interpolate_util.interpolate(
training_example.grid, element_centers, training_example.world2grid)
mse = model_config.hparams.gs * tf.reduce_mean(
tf.sign(gt_sdf_at_centers) * tf.square(gt_sdf_at_centers + 1e-04)) + 1e-5
summarize.summarize_loss(model_config, mse, 'lowres_grid_magnitude_loss')
return mse
def element_center_lowres_grid_inside_loss(model_config, training_example,
structured_implicit):
"""Loss that element centers should lie within a voxel of the GT inside."""
element_centers = structured_implicit.element_centers
gt_sdf_at_centers, _ = interpolate_util.interpolate(
training_example.grid, element_centers, training_example.world2grid)
gt_sdf_at_centers = tf.where_v2(gt_sdf_at_centers > model_config.hparams.igt,
gt_sdf_at_centers, 0.0)
mse = model_config.hparams.ig * tf.reduce_mean(
tf.square(gt_sdf_at_centers + 1e-04)) + 1e-05
summarize.summarize_loss(model_config, mse, 'lowres_grid_inside_loss')
return mse
def overlap_loss(model_config, training_example, structured_implicit):
"""A loss on the overlap between RBF weights."""
sample_count = (model_config.hparams.xsc if model_config.hparams.lrf == 'x'
else model_config.hparams.spc)
samples, _ = training_example.sample_sdf_near_surface(
sample_count=sample_count)
rbf_influences = structured_implicit.rbf_influence_at_samples(samples)
assert len(rbf_influences.shape) == 3 # [b, sample_count, eec]
loss = tf.reduce_mean(tf.linalg.norm(rbf_influences, ord=1,
axis=2)) * model_config.hparams.ow
summarize.summarize_loss(model_config, loss, 'rbf-l1-loss')
return loss
def smooth_element_center_lowres_grid_inside_loss(model_config,
training_example,
structured_implicit):
"""Offset version of element_center_lowres_grid_inside_loss by voxel width."""
element_centers = structured_implicit.element_centers
gt_sdf_at_centers, _ = interpolate_util.interpolate(
training_example.grid, element_centers, training_example.world2grid)
gt_sdf_at_centers = tf.maximum(gt_sdf_at_centers - model_config.hparams.igt,
0.0)
mse = model_config.hparams.ig * tf.reduce_mean(
tf.square(gt_sdf_at_centers + 1e-04)) + 1e-05
summarize.summarize_loss(model_config, mse, 'lowres_grid_inside_loss')
return mse
def center_variance_loss(model_config, training_example, structured_implicit): # pylint:disable=unused-argument
"""A loss on the -variance of the center locations."""
# Training example present for interface uniformity
element_centers = structured_implicit.element_centers
center_shape = element_centers.get_shape().as_list()
if len(center_shape) != 3:
raise ValueError(f'Expected the element centers to have shape [b, #, 3],'
f' but they have shape {center_shape}. center_variance.')
variance = tf.math.reduce_variance(element_centers, axis=[1, 2])
loss_max = model_config.hparams.vt
loss = model_config.hparams.vw * tf.reduce_mean(
tf.maximum(loss_max - variance, 0.0))
summarize.summarize_loss(model_config, loss, 'center-variance-loss')
return loss
def uniform_sample_loss(model_config, training_example, structured_implicit):
"""Loss that uniformly sampled points should have the right insidedness."""
sample_count = (model_config.hparams.xsc if model_config.hparams.lrf == 'x'
else model_config.hparams.spc)
samples, gt_sdf = training_example.sample_sdf_uniform(
sample_count=sample_count)
tf.logging.info('Building Uniform Sample Loss.')
tf.logging.info('Uni. Samples shape: %s',
str(samples.get_shape().as_list()))
loss = model_config.hparams.l2w * sample_loss(model_config,
gt_sdf,
structured_implicit,
samples,
'uniform_sample',
apply_ucf=True)
summarize.summarize_loss(model_config, loss, 'uniform_sample')
return loss
def old_shape_element_center_loss(model_config, training_example,
structured_implicit):
"""Deprecated version of shape_element_center_loss()."""
element_centers = structured_implicit.element_centers
tf.logging.info('Element Center Shape: %s',
str(element_centers.get_shape().as_list()))
bounding_box = training_example.sample_bounding_box
bounding_box_error = tf.reduce_mean(
bounding_box_constraint_error(element_centers, bounding_box))
constraint_loss = model_config.hparams.ibblw * bounding_box_error
summarize.summarize_loss(model_config, constraint_loss,
'inside-bounding-box-loss')
class_at_centers, _ = structured_implicit.class_at_samples(element_centers)
center_loss = tf.reduce_mean((class_at_centers - 0) * (class_at_centers - 0))
center_loss *= model_config.hparams.cclw
summarize.summarize_loss(model_config, center_loss, 'inside-pred-center-loss')
return constraint_loss + center_loss
def near_surface_sample_loss(model_config, training_example,
structured_implicit):
"""An inside/outside loss that samples based on distance to the surface."""
sample_count = (model_config.hparams.xsc if model_config.hparams.lrf == 'x'
else model_config.hparams.spc)
samples, gt_sdf = training_example.sample_sdf_near_surface(
sample_count=sample_count)
tf.logging.info('Building Near Surface Sample Loss.')
tf.logging.info('NS Samples shape: %s', str(samples.get_shape().as_list()))
# TODO(kgenova) Currently we set ucf=True here because that's how it was...
# but go back and fix that because it seems bad.
loss = model_config.hparams.a2w * sample_loss(model_config,
gt_sdf,
structured_implicit,
samples,
'ns_sample',
apply_ucf=True) # False)
summarize.summarize_loss(model_config, loss, 'near_surface_sample')
return loss
def inside_box_loss(model_config, _, structured_implicit):
"""Loss that centers should be inside a fixed size bounding box."""
element_centers = structured_implicit.element_centers
if model_config.hparams.wm == 'f':
bounding_box = shapenet.BoundingBox(lower=-0.7, upper=0.7)
elif model_config.hparams.wm == 't':
bounding_box = shapenet.BoundingBox(
lower=np.array([-.75, -.075, -.75], dtype=np.float32),
upper=np.array([.75, .075, .75], dtype=np.float32))
if model_config.hparams.rsl != 1.0:
bounding_box.lower *= model_config.hparams.rsl
bounding_box.upper *= model_config.hparams.rsl
bounding_box_error = tf.reduce_mean(
bounding_box_constraint_error(element_centers, bounding_box))
outside_bounding_box_loss = model_config.hparams.ibblw * bounding_box_error
summarize.summarize_loss(model_config, outside_bounding_box_loss,
'fixed_bounding_box_loss')
return outside_bounding_box_loss
def shape_element_center_loss(model_config, training_example,
structured_implicit):
"""Loss that centers should be inside the predicted surface."""
element_centers = structured_implicit.element_centers
tf.logging.info('BID0: Shape Element Center Loss.')
tf.logging.info('Element Center Shape: %s',
str(element_centers.get_shape().as_list()))
class_at_centers, _ = structured_implicit.class_at_samples(element_centers)
bounding_box = training_example.sample_bounding_box
bounding_box_error = tf.reduce_mean(
bounding_box_constraint_error(element_centers, bounding_box),
axis=-1,
keep_dims=True)
center_is_inside_gt_box = bounding_box_error <= 0.0
inside_prediction_weights = model_config.hparams.cc * tf.cast(
center_is_inside_gt_box, tf.float32)
# bounding_box_error has shape [batch_size, center_count, 1]
# inside_prediction_weights has shape [batch_size, center_count, 1]
# class_at_centers has shape [batch_size, center_count, 1]. (Double check).
# The class loss is 0 where the prediction is outside the bounding box,
# because the bounding box loss is applied to those centers instead.
class_loss = weighted_l2_loss(0.0, class_at_centers,
inside_prediction_weights)
summarize.summarize_loss(model_config, math_util.nonzero_mean(class_loss),
'ec_loss_class_comp_mean')
outside_bounding_box_loss = model_config.hparams.ibblw * bounding_box_error
summarize.summarize_loss(model_config,
math_util.nonzero_mean(outside_bounding_box_loss),
'ec_loss_outside_bb_comp_mean')
final_loss = tf.reduce_mean(class_loss + outside_bounding_box_loss)
summarize.summarize_loss(model_config, final_loss, 'ec_loss')
return final_loss
def shape_element_center_magnitude_loss(model_config, _, structured_implicit):
element_centers = structured_implicit.element_centers
mse = model_config.hparams.cm * tf.reduce_mean(
tf.square(element_centers + 1e-04)) + 1e-5
summarize.summarize_loss(model_config, mse, 'center_magnitude_loss')
return mse
