simclr-tsne-vis.py

import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import seaborn as sns
tsne = TSNE()

def plot_vecs_n_labels(v,labels,fname):
    fig = plt.figure(figsize = (10, 10))
    plt.axis('off')
    sns.set_style("darkgrid")
    sns.scatterplot(v[:,0], v[:,1], hue=labels, legend='full', palette=sns.color_palette("bright", 5))
    plt.legend(['car', 'dog', 'elephant','cat','airplane'])
    plt.savefig(fname)
    plt.close()
    
for (_, sample_batched) in enumerate(dataloader_training_dataset):
    x = sample_batched['image']
    x = x.to(device)
    y = resnet(x)
    y_tsne = tsne.fit_transform(y.cpu().data)
    labels = sample_batched['label']
    plot_vecs_n_labels(y_tsne,labels,'tsne_train_last_layer.png')