def test_contrastive_loss_computation(self, temperature):
batch_size = 2
project_dim = 16
projection_norm = False
p_1_arr = np.random.rand(batch_size, project_dim)
p_1 = tf.constant(p_1_arr, dtype=tf.float32)
p_2_arr = np.random.rand(batch_size, project_dim)
p_2 = tf.constant(p_2_arr, dtype=tf.float32)
losses_obj = contrastive_losses.ContrastiveLoss(
projection_norm=projection_norm, temperature=temperature)
comp_contrastive_loss = losses_obj(projection1=p_1, projection2=p_2)
def _exp_sim(p1, p2):
return np.exp(np.matmul(p1, p2) / temperature)
l11 = -np.log(
_exp_sim(p_1_arr[0], p_2_arr[0]) /
(_exp_sim(p_1_arr[0], p_1_arr[1]) +
_exp_sim(p_1_arr[0], p_2_arr[1]) +
_exp_sim(p_1_arr[0], p_2_arr[0]))) - np.log(
_exp_sim(p_1_arr[0], p_2_arr[0]) /
(_exp_sim(p_2_arr[0], p_2_arr[1]) +
_exp_sim(p_2_arr[0], p_1_arr[1]) +
_exp_sim(p_1_arr[0], p_2_arr[0])))
l22 = -np.log(
_exp_sim(p_1_arr[1], p_2_arr[1]) /
(_exp_sim(p_1_arr[1], p_1_arr[0]) +
_exp_sim(p_1_arr[1], p_2_arr[0]) +
_exp_sim(p_1_arr[1], p_2_arr[1]))) - np.log(
_exp_sim(p_1_arr[1], p_2_arr[1]) /
(_exp_sim(p_2_arr[1], p_2_arr[0]) +
_exp_sim(p_2_arr[1], p_1_arr[0]) +
_exp_sim(p_1_arr[1], p_2_arr[1])))
exp_contrastive_loss = (l11 + l22) / 2.0
self.assertAlmostEqual(comp_contrastive_loss[0].numpy(),
exp_contrastive_loss,
places=5)
def build_losses(self,
labels,
model_outputs,
aux_losses=None) -> Dict[str, tf.Tensor]:
# Compute contrastive relative loss
con_losses_obj = contrastive_losses.ContrastiveLoss(
projection_norm=self.task_config.loss.projection_norm,
temperature=self.task_config.loss.temperature)
# The projection outputs from model has the size of
# (2 * bsz, project_dim)
projection_outputs = model_outputs[simclr_model.PROJECTION_OUTPUT_KEY]
projection1, projection2 = tf.split(projection_outputs, 2, 0)
contrast_loss, (contrast_logits, contrast_labels) = con_losses_obj(
projection1=projection1, projection2=projection2)
contrast_accuracy = tf.equal(tf.argmax(contrast_labels, axis=1),
tf.argmax(contrast_logits, axis=1))
contrast_accuracy = tf.reduce_mean(
tf.cast(contrast_accuracy, tf.float32))
contrast_prob = tf.nn.softmax(contrast_logits)
contrast_entropy = -tf.reduce_mean(
tf.reduce_sum(contrast_prob * tf.math.log(contrast_prob + 1e-8),
-1))
model_loss = contrast_loss
losses = {
'contrast_loss': contrast_loss,
'contrast_accuracy': contrast_accuracy,
'contrast_entropy': contrast_entropy
}
if self.task_config.model.supervised_head is not None:
outputs = model_outputs[simclr_model.SUPERVISED_OUTPUT_KEY]
labels = tf.concat([labels, labels], 0)
if self.task_config.evaluation.one_hot:
sup_loss = tf.keras.losses.CategoricalCrossentropy(
from_logits=True,
reduction=tf.keras.losses.Reduction.NONE)(labels, outputs)
else:
sup_loss = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True,
reduction=tf.keras.losses.Reduction.NONE)(labels, outputs)
sup_loss = tf.reduce_mean(sup_loss)
label_acc = tf.equal(tf.argmax(labels, axis=1),
tf.argmax(outputs, axis=1))
label_acc = tf.reduce_mean(tf.cast(label_acc, tf.float32))
model_loss = contrast_loss + sup_loss
losses.update({
'accuracy': label_acc,
'supervised_loss': sup_loss,
})
total_loss = model_loss
if aux_losses:
reg_loss = tf.reduce_sum(aux_losses)
total_loss = model_loss + reg_loss
losses['total_loss'] = total_loss
return losses