{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Implement triplet loss and online triplet mining\n\nCross-entropy loss is often used as the training objective of classification tasks. The last layer in the NN model contains K neurons, where K equals to the number of classes we want to classify. \n\n* what if we don't know the total number of classes? i.e. the number of classes keeps varying as new data added into or old data removed from the database daily? \n* only a few samples in each class, i.e. we have 10K images with 1K classes\n* Sometimes we don't actually need to classify data. In face recognition for instance, what we want is to compare two faces and say whether they are from the same person or not. \n\nSiamese neural network or a model trained with triplet loss are often used to learn good embeddings (or “encodings”) that tends to separate images (or other types of data) from different classes.\n\nIn this notebook, we implement a simple CNN model trained with triplet loss. The hardest part of this method is how to select triplets. There are two main triplets mining methods\n* Offline triplet mining: At the beginning of each epoch, compute all the embeddings on the training set, and then only select hard or semi-hard triplets for computing the loss.\n* Online triplet mining: Given a batch of B examples, compute the B embeddings. We then can find a maximum of B^3 triplets.\n\nOnline triplet mining is much more efficient than offline triplet mining. There are two strategies we can choose from while implementing online triplet mining.\n* batch all: select all the valid triplets, and average the loss on the hard and semi-hard triplets\n* batch hard: for each anchor, select the hardest positive and the hardest negative among the batch\n\nFor more details about this topic, please check the references below.\n\n### References\n* (tensorflow implement) https://omoindrot.github.io/triplet-loss#batch-hard-strategy\n* (pytorch implement) https://github.com/NegatioN/OnlineMiningTripletLoss\n* (ranking loss) https://gombru.github.io/2019/04/03/ranking_loss/","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport time\nimport math\nimport random\nimport numpy as np \nimport pandas as pd \nfrom collections import Counter\n\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\nfrom sklearn.manifold import TSNE\nfrom sklearn.model_selection import train_test_split\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import TensorDataset, DataLoader\n\nimport torchvision\nimport torchvision.transforms as T\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### read dataset \n#train = pd.read_csv(\"/kaggle/input/digit-recognizer/train.csv\")\ntrain = pd.read_csv(\"/kaggle/input/fashionmnist/fashion-mnist_train.csv\")\ntrain_x = train[list(train.columns)[1:]].values\ntrain_y = train['label'].values\n\n## normalize and reshape the predictors  \ntrain_x = train_x / 255\n\n## create train and validation datasets\ntrain_x, val_x, train_y, val_y = train_test_split(train_x, train_y, test_size=0.2)\n\n## reshape the inputs\ntrain_x = train_x.reshape(-1, 784)\nval_x = val_x.reshape(-1, 784)\n\n# make dataset \ntrain_dataset = TensorDataset(torch.tensor(train_x), torch.tensor(train_y))\nval_dataset = TensorDataset(torch.tensor(val_x), torch.tensor(val_y))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"X_reduced = TSNE(n_components=2).fit_transform(val_x)\npalette = np.array(sns.color_palette(\"hls\", 10))\nplt.scatter(X_reduced[:, 0], X_reduced[:, 1], c=palette[val_y])\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def my_collate(batch):\n    data = [item[0].view(1, 28,28) for item in batch]\n    data = torch.stack(data).float()\n\n    target = [item[1] for item in batch]\n    target = torch.LongTensor(target)\n    return data, target","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class convEmbedding(nn.Module):\n    def __init__(self, c_dim = 1):\n        super().__init__()\n\n        self.embedding = nn.Sequential(\n            nn.Conv2d(c_dim, 32, 3),\n            nn.ReLU(),\n            nn.Conv2d(32, 32, 3),\n            nn.ReLU(),\n            nn.MaxPool2d(2, 2),\n            nn.Dropout(0.2),\n            \n            nn.Conv2d(32, 64, 3),\n            nn.ReLU(),\n            nn.Conv2d(64, 64, 3),\n            nn.ReLU(),\n            nn.MaxPool2d(2, 2),\n            nn.Dropout(0.2),\n\n            nn.Flatten(),\n\n            nn.Linear(1024, 512),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n\n            nn.Linear(512, 10)\n        )\n\n\n    def forward(self, input_tensor):\n        x = self.embedding(input_tensor)\n        return F.normalize(x, p =2, dim = 1)\n    \nmodel = convEmbedding()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_anchor_positive_mask(labels):\n    \"\"\"Return a 2D mask where mask[a, p] is True iff a and p are distinct and have same label.\n    \"\"\"\n    # Check that i and j are distinct\n    indices_equal = torch.eye(labels.size(0)).bool()\n    indices_not_equal = ~indices_equal\n\n    # Check if labels[i] == labels[j]\n    # Uses broadcasting where the 1st argument has shape (1, batch_size) and the 2nd (batch_size, 1)\n    labels_equal = labels.unsqueeze(0) == labels.unsqueeze(1)\n\n    return labels_equal & indices_not_equal\n\n\ndef get_anchor_negative_mask(labels):\n    \"\"\"Return a 2D mask where mask[a, n] is True iff a and n have distinct labels.\n    \"\"\"\n    # Check if labels[i] != labels[k]\n    # Uses broadcasting where the 1st argument has shape (1, batch_size) and the 2nd (batch_size, 1)\n\n    return ~(labels.unsqueeze(0) == labels.unsqueeze(1))\n\ndef get_triplet_mask(labels):\n    \"\"\"Return a 3D mask where mask[a, p, n] is True iff the triplet (a, p, n) is valid.\n    \"\"\"\n    mask_anchor_positive = get_anchor_positive_mask(labels)\n    mask_anchor_negative = get_anchor_negative_mask(labels)\n\n    return mask_anchor_positive.unsqueeze(2) & mask_anchor_negative.unsqueeze(1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def batch_triplet_loss(embeddings, labels, hparams):\n    \n    pairwise_dist = torch.cdist(embeddings, embeddings, p = 1)\n    \n    if hparams[\"mode\"] == 'batch hard':\n        ##  for each anchor, select the hardest positive and the hardest negative among the batch\n        ## get the hardest positive\n        mask_anchor_positive = get_anchor_positive_mask(labels).float()\n        anchor_positive_dist = pairwise_dist * mask_anchor_positive\n        hardest_positive_dist, _ = anchor_positive_dist.max(1, keepdim=True)\n\n        ## get the hardest negative\n        mask_anchor_negative = get_anchor_negative_mask(labels).float()\n        anchor_negative_dist = pairwise_dist + 999. * (1.0 - mask_anchor_negative)\n        hardest_negative_dist, _ = anchor_negative_dist.min(1, keepdim=True)\n        \n        triplet_loss = hardest_positive_dist - hardest_negative_dist + hparams[\"margin\"] ## (batch_size, 1)\n        \n    elif hparams[\"mode\"] == 'batch all':\n        ## select 1) all the valid triplets, and 2) average the loss on the hard and semi-hard triplets\n        anchor_positive_dist = pairwise_dist.unsqueeze(2) ## (batch_size, batch_size, 1)\n        anchor_negative_dist = pairwise_dist.unsqueeze(1) ## (batch_size, 1, batch_size)\n        triplet_loss = anchor_positive_dist - anchor_negative_dist + hparams[\"margin\"] ## broadcasting, all compibations of (a, p, n)\n        ## 1) select valid triplets\n        valid_mask = get_triplet_mask(labels).float()\n        triplet_loss = valid_mask * triplet_loss ## (batch_size, batch_size, batch_size)\n    else:\n        raise TypeError(\"invalid mode\")\n        \n    # remove easy triplets\n    triplet_loss[triplet_loss < 0] = 0\n    return torch.mean(triplet_loss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def timeSince(since):\n    now = time.time()\n    s = now - since\n    m = math.floor(s / 60)\n    s -= m * 60\n    return '%dm %ds' % (m, s)\n\ndef train_one_batch(model, x, y, hparams):\n    x = x.to(device)\n\n    batch_embeddings = model(x)\n    loss = batch_triplet_loss(batch_embeddings, y, hparams)\n    \n    optimizer.zero_grad()\n    loss.backward()\n    optimizer.step()\n    return loss.item()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"hparams = {\"learning_rate\": 1e-3, \"margin\": 0.5, \"mode\": \"batch hard\", \"batch_size\":  16, \"epochs\": 50} \n\ntrainloader = DataLoader(train_dataset, batch_size = hparams[\"batch_size\"], collate_fn = my_collate, drop_last = True)\nsteps_per_epoch  = train_dataset.__len__()// hparams[\"batch_size\"]\n\nmodel = convEmbedding().to(device)\noptimizer = torch.optim.Adam(model.parameters(), hparams['learning_rate'])\n\n  \nstart = time.time()\ntrain_losses = []\nmodel.train()\nfor epoch in range(1, hparams[\"epochs\"]+1):\n    print('-' * 10)\n    print('Epoch {}/{}\\t{} batches'.format(epoch, hparams[\"epochs\"], steps_per_epoch))\n    \n    curr_loss = []\n    for step, (x, y) in enumerate(trainloader):\n        loss = train_one_batch(model, x, y, hparams)\n        curr_loss.append(loss)\n        print('\\rprogress {:6.1f} %\\tloss {:8.4f}'.format(100*(step+1)/steps_per_epoch, np.mean(curr_loss)), end = \"\")\n        \n    train_losses.append(np.mean(curr_loss))\n    print('\\rprogress {:6.1f} %\\tloss {:8.4f}'.format(100*(step+1)/steps_per_epoch, np.mean(curr_loss)))\n    print('{}'.format(timeSince(start)))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## Check if we can get at least one pair of anchor-positive using this batch size\n#for step, (x, y) in enumerate(trainloader):\n#    print(Counter(y.numpy()))\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\n\nval_loader = DataLoader(val_dataset, batch_size = 128, collate_fn = my_collate)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"X = []; labels = []\nfor x, label in val_loader:\n    batch_embedings = model(x.to(device))\n    X.append(batch_embedings.cpu().detach().numpy())\n    labels.append(label.numpy())\n    print('\\rprogress {:6.1f} %'.format(100*(step+1)/steps_per_epoch), end = \"\")\n\nX = np.concatenate(X)\nlabels = np.concatenate(labels)\n\nX_reduced = TSNE(n_components=2).fit_transform(X)\npalette = np.array(sns.color_palette(\"hls\", 10))\nplt.scatter(X_reduced[:, 0], X_reduced[:, 1], c=palette[labels])\nplt.show()","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}