def test_contrastive_loss(self):
inf_embeddings = tf.nn.l2_normalize(tf.ones((5, 1, 10),
dtype=tf.float32),
axis=-1)
con_embeddings = tf.nn.l2_normalize(tf.ones((5, 1, 10),
dtype=tf.float32),
axis=-1)
tec.compute_embedding_contrastive_loss(inf_embeddings, con_embeddings)
def model_train_fn(
self,
features,
labels,
inference_outputs,
mode,
config = None,
params = None
):
"""Returns weighted sum of losses and individual losses. See base class."""
bc_loss = self._action_decoder.loss(labels)
bc_loss = tf.identity(bc_loss, name='bc_loss')
embed_loss = tec.compute_embedding_contrastive_loss(
inference_outputs['inference_embedding'],
inference_outputs['condition_embedding'])
end_loss = self._compute_end_loss(inference_outputs, labels)
train_outputs = {'bc_loss': bc_loss, 'embed_loss': embed_loss,
'end_loss': end_loss}
return (bc_loss + self._embed_loss_weight * embed_loss +
self._predict_end_weight * end_loss, train_outputs) # pytype: disable=bad-return-type