Contrastive loss with tf.Data

This repository is based on triplet-reid which provides a convenient way to sample images for training. Given a number of classes P and number of samples (images) per class K, triplet-reid leverage tf.data to create a mini-batch of size P*K. This combination, of efficient data loading and specifying the number of classes and samples per classes, is valuable for benchmarking different space embedding approaches.

Unfortunately, tf.contrastive expects embeddings_anchor, embeddings_positive and labels. This doesn't integrate naturally with triplet-reid data loading pipeline. This wiki explains a workaround to integrate these three pieces together, i.e, how to use tf.data with tf.contrastive and control the number of classes and samples per class in a mini-batch.

The following figure illustrates the proposed workaround. For illustration purpose, a mini-batch containing P=2 classes (X, Y) and K=4 samples per classes are utilized. Triplet-reid loads images in order as shown in the next figure. After computing the embeddings, I change the embeddings order by swapping the 3rd and 5th embedding using tf.gather_nd. Finally, I split the embedding into two halves using tf.unstack to prepare the embedding in tf.contrastive format. In this example (P=2, K=4), we end with two pairs of positive and negative embeddings.

The proposed workaround works for P (number of classes) > 2 but has its limitation. It assumes an even number of classes per mini-batch to end with a balanced number of positive and negative embeddings. While K doesn't have to be 4 as in the example, using different K requires code tweaking.