<jupyter_start><jupyter_text><jupyter_code>import numpy as np
import random
import matplotlib.pyplot as plt
from itertools import combinations 
from collections import OrderedDict
from tqdm import tqdm

import tensorflow as tf
 
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
 
import torchvision
import torchvision.transforms as transforms
# plt.rcParams['figure.figsize'] = (5, 5)<jupyter_output><empty_output><jupyter_text>Utils<jupyter_code>def triplet_loss(anchor, positive, negative, margin=0.2):
    pos_dist = (anchor - positive).pow(2).sum(-1) #.pow(.5)
    neg_dist = (anchor - negative).pow(2).sum(-1) #.pow(.5)
    loss = F.relu(pos_dist - neg_dist + margin)
    return loss.mean()
 
# ref: https://github.com/adambielski/siamese-triplet/blob/master/losses.py#L24
def tf_triplet_loss(anchor, positive, negative, margin=0.2):  
    pos_dist = tf.reduce_sum(tf.square(anchor - positive), axis=-1)
    neg_dist = tf.reduce_sum(tf.square(anchor - negative), axis=-1)
    loss = tf.maximum(pos_dist - neg_dist + margin, 0)
    return tf.reduce_mean(loss)
 
# ref : https://stackoverflow.com/questions/38260113/implementing-contrastive-loss-and-triplet-loss-in-tensorflow/38270293#38270293
torch_triplet_loss = nn.TripletMarginLoss(margin=0.2)
def get_all_triplets(embeddings:list, targets:list) -> list:
    '''
    Parameters
    ----------
    embeddings : list of torch.tensor
    targets    : list of ints

    Returns
    -------
    triplets   : list of tuple
    '''

    # eg : no(targets) = 3
    # eg : no(embeds) = 10

    assert len(embeddings) == len(targets) , "Embeddings and Targets must have same lenght"

    triplets = []
    for i, anchor in enumerate(embeddings):

        positive_pairs = list(combinations(anchor,2))  # this will give of distinct pairs of elements 
                                                                 # no(pos_pairs) = no(targets)P2 / 2!  # eg : 45
        
        temp = embeddings.pop(i)  # embeddings list except the anchor
        for negative_embeddings in torch.cat(embeddings,dim=0):   # loops runs for no(targets)-1 * no(embeds) times  # eg : (3-1)*10=20 
            
            triple = [(positive_pair[0], positive_pair[1], negative_embeddings) for positive_pair in positive_pairs]  # no(triple) = n(pos_pair) # eg: 45
            
            triplets.extend(triple)   # no. of triplets added = no(pos_pairs)*(no(targets) - 1)*(no(embeds))  # eg: 45*2*10=900

        embeddings.insert(i, temp)

    return triplets  # no(triplets) = no(added_triples) * no(embeds) # eg: 900*3 = 2700

# ref : A lot can't even even mention here
def get_all_tensor_triplets(embeddings:list, targets:list) -> torch.Tensor:
    '''
    Parameters
    ----------
    embeddings : list of torch.tensor each of shape torch.Size([?, 128])
    targets    : list of ints

    Returns
    -------
    triplets   : torch.tensor of shape torch.Size([?, no(triplets), 3])
    '''

    # eg : no(targets) = 3
    # eg : no(embeds) = 10

    assert len(embeddings) == len(targets) , "Embeddings and Targets must have same lenght"

    triplets = []
    for i, anchor in enumerate(embeddings):

        positive_pairs = list(combinations(anchor,2))  # this will give of distinct pairs of elements 
                                                                 # no(pos_pairs) = no(targets)P2 / 2!  # eg : 45
        
        temp = embeddings.pop(i)  # embeddings list except the anchor
        for negative_embeddings in torch.cat(embeddings,dim=0):   # loops runs for no(targets)-1 * no(embeds) times  # eg : (3-1)*10=20 
            
            triple = [torch.stack([positive_pair[0], positive_pair[1], negative_embeddings], dim=0) for positive_pair in positive_pairs]  # no(triple) = n(pos_pair) # eg: 45
            
            triplets.extend(triple)   # no. of triplets added = no(pos_pairs)*(no(targets) - 1)*(no(embeds))  # eg: 45*2*10=900

        embeddings.insert(i, temp)

    return torch.stack(triplets, dim=0)  # no(triplets) = no(added_triples) * no(embeds) # eg: 900*3 = 2700

# ref : A lot can't even even mention here
def get_random_triplets(embeddings:list,  targets:list) -> torch.Tensor:
    '''
    For each image in data (Anchor), randomly sample a Positive image from its class.  
    Then from each of the other classes sample one image (randomly) as the Negative. Hence, for every Anchor
    you will have 1 randomly selected positive from it's class and randomly selected Negative from each of the n-1 classes 
    where n is total number of classes.  For every Anchor image you would have n-1 triplets.
    So if you're having 3 classes of 10 images each then you would have 60 triplets. 
    '''


    assert len(embeddings) == len(targets) , "Embeddings and Targets must have same lenght"

    triplets = []
    for i, embedding in enumerate(embeddings):
        temp = embeddings.pop(i)

        for anchor in embedding:
            positive = random.choice(embedding)

            for negatives in embeddings:
                negative = random.choice(negatives)

                triplets.append(torch.stack([anchor, positive, negative], dim=0))

        embeddings.insert(i, temp)

    return torch.stack(triplets, dim=0)  

colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728',
              '#9467bd', '#8c564b', '#e377c2', '#7f7f7f',
              '#bcbd22', '#17becf']
 
def plot_embeddings(embeddings, targets, xlim=None, ylim=None):
    plt.figure(figsize=(10,10))
    for i in range(3):
        plt.scatter(embeddings[i].detach().numpy()[:,0],embeddings[i].detach().numpy()[:,1], alpha=0.5, color=colors[i])
    if xlim:
        plt.xlim(xlim[0], xlim[1])
    if ylim:
        plt.ylim(ylim[0], ylim[1])
    plt.legend(targets)

# ref: https://github.com/adambielski/siamese-triplet/blob/master/Experiments_MNIST.ipynb
MNIST_train = torchvision.datasets.MNIST(
    root = '/content/drive/MyDrive/Colab Notebooks/',
    download = True,
    train = True,
    transform = transforms.ToTensor()
)
MNIST_train.data.shape
i =5734
plt.imshow(MNIST_train[i][0].squeeze())
MNIST_train[i][1]<jupyter_output><empty_output><jupyter_text>Reorganizing Data<jupyter_code>train_dict = {i : MNIST_train.data[MNIST_train.targets == i].unsqueeze(1) / 255 for i in range(10)}
print(*[f'{imgs.shape[0]} images of Label {label} of shape {imgs.shape[1:]}' for label, imgs in train_dict.items()], sep='\n')
# for i in range(10):
#     train_dict[i] = MNIST_train.data[MNIST.targets == i]
label = 7
rand_imgs = train_dict[label][random.sample(range(1, 5000), 100)]
grid = torchvision.utils.make_grid(rand_imgs, nrow=10, pad_value=20)
# ref: https://pytorch.org/vision/stable/utils.html#torchvision-utils
plt.figure(figsize=(15,15))
plt.imshow(np.transpose(grid, (1,2,0)))<jupyter_output>Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).<jupyter_text>Image Processing<jupyter_code>preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    transforms.Normalize(mean=0.1307, std=0.3081),
    transforms.Lambda(lambda image: torch.cat((image,image,image), dim=-3)),
])
# ref: https://pytorch.org/hub/pytorch_vision_resnet/
timg = train_dict[5][63]  # shape = (1,28,28)
ptimg = preprocess(timg)  # shape = (3,224,224)
_, axarr = plt.subplots(1,2, figsize = (7,7)) 

axarr[0].imshow(timg.squeeze())

axarr[1].imshow(torch.movedim(ptimg, -3, -1))
#ref : https://pytorch.org/docs/stable/generated/torch.movedim.html
rand_imgs = train_dict[1][random.sample(range(1, 5000), 5)]  # shape = (5,1,28,28)
grid = torchvision.utils.make_grid(rand_imgs, nrow=5, pad_value=20)

prand_imgs = preprocess(rand_imgs)  # shape = (5,3,224,224)
pgrid = torchvision.utils.make_grid(prand_imgs, nrow=5, pad_value=20)

fig , axarr = plt.subplots(2,1, figsize = (15,5)) 

axarr[0].imshow(np.transpose(grid, (1,2,0)))
axarr[1].imshow(np.transpose(pgrid, (1,2,0)))
IMAGES = {0:train_dict[0][95], 1:train_dict[1][1], 2:train_dict[2][2], 3:train_dict[3][0]}
for i in [4,6,7,8]:
    IMAGES[i] = train_dict[i][1]
IMAGES[5] = train_dict[5][63]  
IMAGES[9] = train_dict[9][5]
GRID = torchvision.utils.make_grid(preprocess(torch.stack([*IMAGES.values()])), nrow=5, pad_value=2)
plt.figure(figsize=(15,5))
plt.imshow(np.transpose(GRID, (1,2,0)))<jupyter_output>Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).<jupyter_text>Creating Model<jupyter_code># Self defined CNN Model


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.preprocess = transforms.Compose([
                            transforms.Resize(256),
                            transforms.CenterCrop(224),
                            # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
                            transforms.Normalize(mean=0.1307, std=0.3081),
                            transforms.Lambda(lambda image: torch.cat((image,image,image), dim=-3)),
                        ])
        # ref: https://pytorch.org/hub/pytorch_vision_resnet/

        self.convnet1 = nn.Sequential(OrderedDict([
                            ('conv1', nn.Conv2d(in_channels=3, out_channels=64, kernel_size=7, stride=2)),
                            ('conv2', nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3)),
                            ('relu', nn.ReLU()),
                            ('pool', nn.MaxPool2d(kernel_size=3, stride=2))
        ]))
        # ref : https://discuss.pytorch.org/t/is-it-possible-to-specify-a-name-for-each-layer-when-creating-the-model/33637/2
        self.convnet2 = nn.Sequential(
                            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3),
                            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3),
                            nn.ReLU(),
                            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.convnet3 = nn.Sequential(
                            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2),
                            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3),
                            nn.ReLU(),
                            nn.MaxPool2d(kernel_size=3, stride=2)
        )
        self.FCs      = nn.Sequential(
                            nn.Linear(in_features=4096, out_features=2048), nn.ReLU(),
                            nn.Linear(in_features=2048, out_features=512), nn.ReLU(),
                            nn.Linear(in_features=512, out_features=128), nn.ReLU()
        )
        self.out      = nn.Linear(in_features=128, out_features=2)


    def forward(self, x):
        x = self.preprocess(x)
        x = self.convnet1(x)
        x = self.convnet2(x)
        x = self.convnet3(x)
        x = x.reshape(-1, 4*4*256)
        x = self.FCs(x)
        enc = self.out(x)
        return x<jupyter_output><empty_output><jupyter_text>Visualizing Model<jupyter_code>model = Model()
model
!pip install torchviz
from torchviz import make_dot

x = torch.randn(10,1,28,28)

make_dot(model(x), params=dict(model.named_parameters()))   

# ref : https://github.com/szagoruyko/pytorchviz/blob/master/examples.ipynb
!pip install git+https://github.com/waleedka/hiddenlayer.git
import hiddenlayer as hl

# Build HiddenLayer graph
hl_graph = hl.build_graph(model, torch.zeros([1, 1, 28, 28]))

# Use a different color theme
hl_graph.theme = hl.graph.THEMES["blue"].copy()  # Two options: basic and blue
hl_graph

#ref : https://github.com/waleedka/hiddenlayer/blob/master/demos/pytorch_graph.ipynb
colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728',
              '#9467bd', '#8c564b', '#e377c2', '#7f7f7f',
              '#bcbd22', '#17becf']
 
def plot_embeddings(embeddings, targets, xlim=None, ylim=None):
    # out = model.out
    plt.figure(figsize=(10,10))
    for i in range(10):
        # emb = out(embeddings[i])
        # plt.scatter(emb.detach().numpy()[:,0],emb.detach().numpy()[:,1], alpha=0.5, color=colors[i])
        plt.scatter(embeddings[i].detach().numpy()[:,0],embeddings[i].detach().numpy()[:,1], alpha=0.5, color=colors[i])
    if xlim:
        plt.xlim(xlim[0], xlim[1])
    if ylim:
        plt.ylim(ylim[0], ylim[1])
    plt.legend(targets)

# ref: https://github.com/adambielski/siamese-triplet/blob/master/Experiments_MNIST.ipynb<jupyter_output><empty_output><jupyter_text>Running the Model<jupyter_code>def weights_init(layer):
    classname = layer.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.kaiming_normal_(layer.weight.data, nonlinearity='relu')  # He-initialization
        # ref : https://pytorch.org/docs/stable/nn.init.html#torch.nn.init.kaiming_normal_
        
# Reference : https://discuss.pytorch.org/t/weight-initilzation/157
model = Model()
model.apply(weights_init)

learning_rate = 0.001
lossfn = triplet_loss
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)

zeros_images = train_dict[0][:100]
ones_images  = train_dict[1][:100]
twos_images  = train_dict[2][:100]

images = torch.stack([zeros_images, ones_images, twos_images], dim=1)
images.shape  # images[:,0,...].shape = torch.Size([100, 1, 28, 28])
loader = torch.utils.data.DataLoader(images, batch_size=10)
len(loader)
# sample = next(iter(loader))
# sample.shape = torch.Size([10, 3, 1, 28, 28])
enc0 = []
enc1 = []
enc2 = []
for batch in tqdm(loader):
    enc0.append(model(batch[:,0,...]))
    enc1.append(model(batch[:,1,...]))
    enc2.append(model(batch[:,2,...]))
enc0 = torch.cat(enc0)
enc1 = torch.cat(enc1)
enc2 = torch.cat(enc2)
enc0.shape, enc1.shape, enc2.shape
plot_embeddings(embeddings=[enc0, enc1, enc2], targets=[0,1,2])
cost = []

epochs = 5
for epoch in range(epochs):
    for batch in loader:
        zeros_encodings = model(batch[:,0,...])
        ones_encodings  = model(batch[:,1,...])
        twos_encodings  = model(batch[:,2,...])

        triplets = get_random_triplets([zeros_encodings, ones_encodings, twos_encodings], [0, 1, 2])

        anchor   = triplets[:,0]
        positive = triplets[:,1]
        negative = triplets[:,2]

        optimizer.zero_grad()
        loss = lossfn(anchor, positive, negative, margin=1)
        cost.append(loss)

        
        loss.backward()
        optimizer.step()

        print(f'Epoch:[{epoch+1}/{epochs}] , Cost : {loss.item()}')

plt.xlabel('Total Iterations')
plt.ylabel('Cost')
plt.plot(np.arange(epochs*len(loader)), cost)
plt.show()
enc0 = []
enc1 = []
enc2 = []
for batch in tqdm(loader):
    enc0.append(model(batch[:,0,...]))
    enc1.append(model(batch[:,1,...]))
    enc2.append(model(batch[:,2,...]))
enc0 = torch.cat(enc0)
enc1 = torch.cat(enc1)
enc2 = torch.cat(enc2)
plot_embeddings([enc0, enc1, enc2], [0,1,2])
# Saving the state_dict of Model
PATH = '/content/drive/MyDrive/Colab Notebooks/tripletLoss_on_MNIST'
# torch.save(model.state_dict(), '/content/Triplet Loss/tripletLoss_on_MNIST')
# Loadel the saved Model
# the_network = Network()
model.load_state_dict(torch.load(PATH))
model.eval()
IMAGES_ENC = {k:model(v.unsqueeze(0)) for k,v in IMAGES.items()}
train_dict[0].shape
zeros = model(train_dict[0][:100])
(zeros - IMAGES_ENC[0]).shape
zero_dist = (zeros - IMAGES_ENC[0]).pow(2).sum(-1)
zero_dist.max() - zero_dist.mean()
correct = {i:0 for i in range(10)}
for k , value in train_dict.items():
    vload = torch.utils.data.DataLoader(value, batch_size=50)
    for v in tqdm(vload, desc=f"Class : {k}", position=0, leave=True):
        correct[k] += ((model(v) - IMAGES_ENC[k]).pow(2).sum(-1) < 12.6).sum()
for i in range(10):
    c = correct[i].item()
    t = train_dict[i].shape[0]
    print(f'Class {i} : Correct = {c} out of {t} . Accuracy = {c / t * 100} %')

imgs = [train_dict[i][:50] for i in range(10)]
imgs = torch.stack(imgs, dim=1)
imgs.shape  # images[:,0,...].shape = torch.Size([100, 1, 28, 28])
loader = torch.utils.data.DataLoader(imgs, batch_size=10)
print(len(loader))
sample = next(iter(loader))
sample.shape #= torch.Size([10, 3, 1, 28, 28])
# enc = {i:[] for i in range(10)}
for batch in tqdm(loader):
    for i in range(10):
        e = model(batch[:,i,...])
        enc[i].append(model.out(e))
    
for i in range(10):
    enc[i] = torch.cat(enc[i])
    print(enc[i].shape)
plot_embeddings([*enc.values()], list(range(10)))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15545963035019456, 'fraction_numerical': 0.04383511673151751, 'mean_word_length': 3.1695817490494296, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 30, 'https://': 14, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '114524', 'license': 'MIT', 'n_tokens_mistral': 6300, 'n_tokens_neox': 5886, 'n_words': 1145, 'path': 'triplet_network/Triplet Loss MNIST/TRIPLET_LOSS_Pytorch.ipynb', 'repo_name': 'IvLabs/Face-Unlock', 'size': '16448'}