{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:19:20.692611Z",
     "start_time": "2020-10-24T06:19:20.621531Z"
    }
   },
   "outputs": [],
   "source": [
    "# default_exp losses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Losses\n",
    "\n",
    "> Custom loss functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:19:21.036103Z",
     "start_time": "2020-10-24T06:19:20.988699Z"
    }
   },
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:19:21.834418Z",
     "start_time": "2020-10-24T06:19:21.751155Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:19:31.247423Z",
     "start_time": "2020-10-24T06:19:31.187317Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "from functools import partial\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from htools import add_docstring, valuecheck, identity\n",
    "from incendio.layers import SmoothLogSoftmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-23T06:05:42.139333Z",
     "start_time": "2020-10-23T06:05:41.768042Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hmamin/anaconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject\n",
      "  return f(*args, **kwds)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject\n",
      "  return f(*args, **kwds)\n"
     ]
    }
   ],
   "source": [
    "# Used for testing only.\n",
    "from fastai2.metrics import LabelSmoothingCrossEntropy\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "\n",
    "from htools import assert_raises, InvalidArgumentError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def smooth_soft_labels(labels, alpha=.1):\n",
    "    \"\"\"Add uniform probability to a tensor of soft labels or \n",
    "    one-hot-encoded hard labels. Classes with zero probability\n",
    "    will therefore end up with alpha/k probability, where k is\n",
    "    the total number of classes. Note that some implementations\n",
    "    may use alpha/(k-1), so if you notice a difference in output \n",
    "    that could be the source.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    labels: torch.tensor\n",
    "        Tensor of labels where each row is a sample and each column\n",
    "        is a class. These can be one hot encodings or a vector of soft\n",
    "        probabilities. Shape (batch_size, num_classes).\n",
    "    alpha: float\n",
    "        A positive value which will be used to assign nonzero \n",
    "        probabilities to the classes that are currently zeros. A larger\n",
    "        alpha corresponds to a higher degree of smoothing (useful when\n",
    "        accounting for noisier labels, trying to provide a stronger\n",
    "        regularizing effect, or encouraging less confident predictions).\n",
    "    \"\"\"\n",
    "    if alpha < 0:\n",
    "        raise InvalidArgumentError('Alpha must be non-negative.')\n",
    "    \n",
    "    # Avoid unnecessary computation.\n",
    "    if not alpha: \n",
    "        return labels\n",
    "    length = labels.shape[-1]\n",
    "    nonzeros = (labels > 0).sum(-1).unsqueeze(-1).float()\n",
    "    return torch.clamp_min(labels - alpha/nonzeros, 0) + alpha/length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def soft_label_cross_entropy_with_logits(y_pred, y_true, alpha=0.0, \n",
    "                                         reduction='mean'):\n",
    "    \"\"\"Compute cross entropy with soft labels. PyTorch's built in \n",
    "    multiclass cross entropy functions require us to pass in integer\n",
    "    indices, which doesn't allow for soft labels which are shaped like\n",
    "    a one hot encoding. FastAI's label smoothing loss uniformly divides\n",
    "    uncertainty over all classes, which again does not allow us to pass\n",
    "    in our own soft labels.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y_pred: torch.FloatTensor\n",
    "        Logits output by the model. \n",
    "        Shape (bs, num_classes).\n",
    "    y_true: torch.FloatTensor\n",
    "        Soft labels, where values are between 0 and 1. \n",
    "        Shape (bs, num_classes).\n",
    "    alpha: float\n",
    "        Label smoothing hyperparameter: a positive value which will be used to\n",
    "        assign nonzero probabilities to the classes that are currently zeros. \n",
    "        A larger alpha corresponds to a higher degree of smoothing (useful when\n",
    "        accounting for noisier labels, trying to provide a stronger\n",
    "        regularizing effect, or encouraging less confident predictions).\n",
    "    reduction: str\n",
    "        One of ('mean', 'sum', 'none'). This determines how to reduce\n",
    "        the output of the function, similar to most PyTorch\n",
    "        loss functions.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    torch.FloatTensor: If reduction is 'none', this will have shape\n",
    "        (bs, ). If 'mean' or 'sum', this will be be a tensor with a \n",
    "        single value (no shape).\n",
    "    \"\"\"\n",
    "    res = (-smooth_soft_labels(y_true, alpha) \n",
    "           * F.log_softmax(y_pred, dim=-1)).sum(-1)\n",
    "    if reduction == 'none': return res\n",
    "    return getattr(res, reduction)(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@add_docstring(soft_label_cross_entropy_with_logits)\n",
    "def soft_label_cross_entropy(y_pred, y_true, alpha=0.0, reduction='mean'):\n",
    "    \"\"\"Same as `soft_label_cross_entropy_with_logits` but operates on\n",
    "    softmax output instead of logits. The version with logits is \n",
    "    recommended for numerical stability. Below is the docstring for the logits \n",
    "    version. The only difference in this version is that y_pred will not be\n",
    "    logits.\n",
    "    \"\"\"\n",
    "    res = -smooth_soft_labels(y_true, alpha) * torch.log(y_pred)\n",
    "    res = torch.where(torch.isnan(res) | torch.isinf(res),\n",
    "                      torch.zeros_like(res), res).sum(-1)\n",
    "    if reduction == 'none': return res\n",
    "    return getattr(res, reduction)(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For demonstration, each row of our soft labels tries to capture a different case: \n",
    "- Row 0: High confidence label for one class, model has high confidence in the correct class.\n",
    "- Row 1: Moderate confidence in 2 different classes, model has high confidence in the best class.\n",
    "- Row 2: Confidence is split between a few classes, model predictions assign most probability to the two nonzero but non-ideal classes.\n",
    "\n",
    "Row 2 should benefit slightly from label smoothing since it predicts answers that are not ideal but not entirely wrong."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[9.2457e-01, 6.2297e-03, 4.6032e-02, 6.2297e-03, 1.6934e-02],\n",
       "        [4.7232e-02, 9.4869e-01, 8.6509e-04, 8.6509e-04, 2.3516e-03],\n",
       "        [3.1685e-02, 2.3412e-01, 6.3641e-01, 8.6129e-02, 1.1656e-02]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 3 row mini batch\n",
    "y_label = torch.tensor([0, 1, 3])\n",
    "y_ohe = torch.tensor([[1, 0, 0, 0, 0],\n",
    "                      [0, 1, 0, 0, 0], \n",
    "                      [0, 0, 0, 1, 0]], dtype=torch.float)\n",
    "y_smooth = torch.tensor([[.9, .1, 0, 0, 0],\n",
    "                         [.35, .65, 0, 0, 0],\n",
    "                         [0, .1, .2, .7, 0]])\n",
    "logits = torch.tensor([[6, 1, 3, 1, 2], \n",
    "                       [5, 8, 1, 1, 2],\n",
    "                       [1, 3, 4, 2, 0]], dtype=torch.float)\n",
    "y_hat = F.softmax(logits, dim=-1)\n",
    "y_hat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got InvalidArgumentError(Alpha must be non-negative.).\n"
     ]
    }
   ],
   "source": [
    "# Test: smooth_soft_labels\n",
    "assert torch.allclose(\n",
    "    smooth_soft_labels(y_smooth, .1), \n",
    "    torch.tensor([[0.8700, 0.0700, 0.0200, 0.0200, 0.0200],\n",
    "                  [0.3200, 0.6200, 0.0200, 0.0200, 0.0200],\n",
    "                  [0.0200, 0.0867, 0.1867, 0.6867, 0.0200]]),\n",
    "atol=1e-3), 'Case 1: alpha=0.1'\n",
    "\n",
    "assert torch.allclose(smooth_soft_labels(y_smooth, 0.0), y_smooth, atol=1e-3),\\\n",
    "    'Case 2: alpha=0.0'\n",
    "\n",
    "with assert_raises(InvalidArgumentError):\n",
    "    smooth_soft_labels(y_smooth, -.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reduction=none: tensor([0.0784, 0.0527, 2.4519])\n",
      "reduction=sum: tensor(2.5830)\n",
      "reduction=mean: tensor(0.8610)\n"
     ]
    }
   ],
   "source": [
    "modes = ('none', 'sum', 'mean')\n",
    "for mode in modes:\n",
    "    j = soft_label_cross_entropy_with_logits(logits, y_ohe, reduction=mode)\n",
    "    print(f'reduction={mode}:', j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "soft_label_cross_entropy_with_logits:\n",
      "-------------------------------------\n",
      "reduction=none: tensor([0.5784, 1.1027, 1.9519])\n",
      "reduction=sum: tensor(3.6330)\n",
      "reduction=mean: tensor(1.2110)\n",
      "\n",
      "soft_label_cross_entropy:\n",
      "-------------------------\n",
      "reduction=none: tensor([0.5784, 1.1027, 1.9519])\n",
      "reduction=sum: tensor(3.6330)\n",
      "reduction=mean: tensor(1.2110)\n"
     ]
    }
   ],
   "source": [
    "print('soft_label_cross_entropy_with_logits:\\n' + '-'*37)\n",
    "for mode in modes:\n",
    "    j = soft_label_cross_entropy_with_logits(logits, y_smooth, reduction=mode)\n",
    "    print(f'reduction={mode}:', j)\n",
    "\n",
    "print('\\nsoft_label_cross_entropy:\\n' + '-'*25)\n",
    "for mode in modes:\n",
    "    j = soft_label_cross_entropy(y_hat, y_smooth, reduction=mode)\n",
    "    print(f'reduction={mode}:', j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reduction=none: tensor([0.7584, 1.7227, 2.1519])\n",
      "reduction=sum: tensor(4.6330)\n",
      "reduction=mean: tensor(1.5443)\n"
     ]
    }
   ],
   "source": [
    "for mode in modes:\n",
    "    j = soft_label_cross_entropy_with_logits(logits, y_smooth, alpha=.2, \n",
    "                                             reduction=mode)\n",
    "    print(f'reduction={mode}:', j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_loss_by_alpha(y_pred, y_true, loss_func=None, loss_class=None):\n",
    "    alphas = np.arange(0, 1, .05)\n",
    "    \n",
    "    if loss_class:\n",
    "        js = {a: loss_class(a, 'none')(y_pred, y_true) for a in alphas}\n",
    "    else:\n",
    "        js = {a: loss_func(y_pred, y_true, a, 'none') for a in alphas}\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    for i in range(3):\n",
    "        ax.plot(alphas, [x[i] for x in js.values()], label=i)\n",
    "    plt.xlabel('Alpha (Smoothing Parameter)')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice the gap between row 2 begins to narrow, as predicted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21668 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_loss_by_alpha(y_hat, y_smooth, soft_label_cross_entropy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21668 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_loss_by_alpha(logits, y_smooth, soft_label_cross_entropy_with_logits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15108 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_loss_by_alpha(logits, y_ohe, soft_label_cross_entropy_with_logits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For comparison, here is the fastai label smoothing loss. It should return the same results in this case (just adding uniform noise). The advantage of ours is the ability to specify a prior non-uniform distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15108 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_loss_by_alpha(logits, y_label, loss_class=LabelSmoothingCrossEntropy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def classification_mae_with_logits(y_pred, y_true, reduction='mean'):\n",
    "    \"\"\"Mean absolute error for classification using soft or one hot encoded\n",
    "    labels (can be helpful for dealing with noisy labels). Classification\n",
    "    models often forego the final softmax layer and use a loss function which\n",
    "    computes the log softmax for numerical stability reasons. This loss \n",
    "    function makes it easy to swap different loss functions in and out without\n",
    "    changing the model or the training loop.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    y_pred: torch.FloatTensor\n",
    "        Logits. Shape (bs, num_classes).\n",
    "    y_true: torch.FloatTensor\n",
    "        Labels, either one hot encoded or soft. Shape (bs, num_classes).\n",
    "    reduction: str\n",
    "        Just like any PyTorch loss function, this determines how to aggregate\n",
    "        the loss over a mini batch. One of ('mean', 'sum', 'none').\n",
    "    \"\"\"\n",
    "    # Built-in reduction will aggregate over every value in the tensor because\n",
    "    # it expects that we're predicting a single value per sample. Because \n",
    "    # we're predicting a vector for each sample, we compute row-wise sums \n",
    "    # regardless of our choice of reduction, then aggregate over the batch \n",
    "    # dimension if specified.\n",
    "    rowwise_j = F.l1_loss(F.softmax(y_pred, dim=-1),\n",
    "                          y_true, reduction='none').sum(-1)\n",
    "    if reduction == 'none': return rowwise_j\n",
    "    return getattr(rowwise_j, reduction)(dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[6., 1., 3., 1., 2.],\n",
       "        [5., 8., 1., 1., 2.],\n",
       "        [1., 3., 4., 2., 0.]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[9.2457e-01, 6.2297e-03, 4.6032e-02, 6.2297e-03, 1.6934e-02],\n",
       "        [4.7232e-02, 9.4869e-01, 8.6509e-04, 8.6509e-04, 2.3516e-03],\n",
       "        [3.1685e-02, 2.3412e-01, 6.3641e-01, 8.6129e-02, 1.1656e-02]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.softmax(logits, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 0., 0.],\n",
       "        [0., 1., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 0.]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_ohe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "none tensor([0.1509, 0.1026, 1.8277])\n",
      "sum tensor(2.0812)\n",
      "mean tensor(0.6937)\n"
     ]
    }
   ],
   "source": [
    "for mode in modes:\n",
    "    j = classification_mae_with_logits(logits, y_ohe, mode)\n",
    "    print(mode, j) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-21T06:02:35.186670Z",
     "start_time": "2020-10-21T06:02:35.152902Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class PairwiseLossReduction(nn.Module):\n",
    "    \"\"\"Basically lets us use L2 or L1 distance as a loss function with the\n",
    "    standard reductions. If we don't want to reduce, we could use the built-in\n",
    "    torch function, but that will usually output a tensor rather than a \n",
    "    scalar.\n",
    "    \"\"\"\n",
    "\n",
    "    @valuecheck\n",
    "    def __init__(self, reduce:('sum', 'mean', 'none')='mean', **kwargs):\n",
    "        super().__init__()\n",
    "        self.distance = nn.PairwiseDistance(**kwargs)\n",
    "        self.reduce = identity if reduce == 'none' else getattr(torch, reduce)\n",
    "\n",
    "    def forward(self, y_proba, y_true):\n",
    "        return self.reduce(self.distance(y_proba, y_true))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:19:55.287107Z",
     "start_time": "2020-10-24T06:19:55.248885Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def reduce(x, reduction='mean'):\n",
    "    \"\"\"This is a very common line in my custom loss functions so I'm providing\n",
    "    a convenience function for it. I think this function call is also more \n",
    "    intuitive than the code behind it.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: torch.Tensor\n",
    "        The object to reduce.\n",
    "    reduction: str\n",
    "        Should be one of ('mean', 'sum', 'none'), though technically some \n",
    "        other operations (e.g. 'std') are supported.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    torch.Tensor: Scalar if using 'mean' or 'sum', otherwise the same as \n",
    "    input `x`.\n",
    "    \n",
    "    Examples\n",
    "    --------\n",
    "    def squared_error(x, reduction='mean'):\n",
    "        return reduce(x.pow(2), reduction)\n",
    "    \"\"\"\n",
    "    return x if reduction == 'none' else getattr(torch, reduction)(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:28:49.480830Z",
     "start_time": "2020-10-24T06:28:49.337174Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def contrastive_loss(x1, x2, y, m=1.25, p=2, reduction='mean'):\n",
    "    \"\"\"Functional form of the contrastive loss as described by Hadsell, \n",
    "    Chopra, and LeCun in \n",
    "    \"Dimensionality Reduction by Learning an Invariant Mapping\":\n",
    "    http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf\n",
    "    \n",
    "    Similar examples always contribute to the loss while negatives only do\n",
    "    if they are sufficiently similar.\n",
    "    \n",
    "    A couple words of caution: this uses the convention that y=1 represents a \n",
    "    positive label (similar examples) despite the paper doing the opposite.\n",
    "    Also, the value of m (margin) might benefit from some tuning as I'm not\n",
    "    very confident in the default value.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x1: torch.Tensor\n",
    "        Shape (bs, n_features).\n",
    "    x2: torch.Tensor\n",
    "        Shape (bs, n_features).\n",
    "    y: torch.Tensor\n",
    "        Labels. Unlike the paper, we use the convention that a label of 1 \n",
    "        means images are similar. This is consistent with all our existing\n",
    "        datasets and just feels more intuitive.\n",
    "    m: float\n",
    "        Margin that prevents dissimilar pairs from affecting the loss unless\n",
    "        they are sufficiently far apart. I believe the reasonable range of\n",
    "        values depends on the size of the feature dimension. The default is\n",
    "        based on a figure in the paper linked above but I'm not sure how\n",
    "        much stock to put in that.\n",
    "    p: int\n",
    "        The p that determines the p-norm used to calculate the initial \n",
    "        distance measure between x1 and x2. The default of 2 therefore uses\n",
    "        euclidean distance.\n",
    "    reduction: str\n",
    "        One of ('sum', 'mean', 'none'). Standard pytorch loss reduction. Keep\n",
    "        in mind 'none' will probably not allow backpropagation since it\n",
    "        returns a rank 2 tensor.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.Tensor: Scalar measuring the contrastive loss. If no reduction is\n",
    "    applied, this will instead be a tensor of shape (bs,).\n",
    "    \"\"\"\n",
    "    dw = F.pairwise_distance(x1, x2, p, keepdim=True) \n",
    "    # Loss_similar + Loss_different\n",
    "    res = y*dw.pow(p).div(2) + (1-y)*torch.clamp_min(m-dw, 0).pow(p).div(2)\n",
    "    return reduce(res, reduction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:28:50.681226Z",
     "start_time": "2020-10-24T06:28:50.587684Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class ContrastiveLoss1d(nn.Module):\n",
    "    \n",
    "    @add_docstring(contrastive_loss)\n",
    "    def __init__(self, m=1.25, p=2, reduction='mean'):\n",
    "        \"\"\"OOP version of contrastive loss. The docs for the functional \n",
    "        version are below:\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.m = m\n",
    "        self.p = p\n",
    "        self.reduction = reduction\n",
    "        self.loss = partial(contrastive_loss, m=m, p=p, reduction=reduction)\n",
    "        \n",
    "    def forward(self, x1, x2, y_true):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x1: torch.Tensor\n",
    "            Shape (bs, feature_dim).\n",
    "        x2: torch.Tensor\n",
    "            Shape (bs, feature_dim).\n",
    "        y_true: torch.Tensor\n",
    "            Shape (bs, 1). 1's indicate the inputs are \"similar\", 0's indicate\n",
    "            they are dissimilar.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor: scalar if `reduction` is not 'none', otherwise tensor \n",
    "        has same shape as `y_true`.\n",
    "        \"\"\"\n",
    "        assert y_true.ndim == 2, \"y_true must be rank 2.\"\n",
    "        return self.loss(x1, x2, y_true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:24:28.730190Z",
     "start_time": "2020-10-24T06:24:28.683808Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ -1.1935,  19.9629,  39.9642,  60.6310,  79.9442],\n",
      "        [-48.8698, -49.4237, -48.8313, -49.1232, -49.0722]])\n",
      "tensor([[-0.8069, 18.7798, 38.6250, 59.6368, 80.2849],\n",
      "        [24.5225, 25.7855, 25.0192, 25.8827, 23.5852]])\n"
     ]
    }
   ],
   "source": [
    "bs = 2\n",
    "x1 = torch.randn(bs, 5)\n",
    "x2 = torch.randn(bs, 5)\n",
    "y = torch.tensor([1, 0]).unsqueeze(-1)\n",
    "\n",
    "# Make pair 1 similar, pair 2 dissimilar.\n",
    "x1[0] += torch.arange(0, 100, 20)\n",
    "x1[1] -= 50\n",
    "x2[0] += torch.arange(0, 100, 20)\n",
    "x2[1] += 25\n",
    "\n",
    "print(x1)\n",
    "print(x2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:24:46.245131Z",
     "start_time": "2020-10-24T06:24:46.184418Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.1118)"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = ContrastiveLoss1d(reduction='mean')\n",
    "loss(x1, x2, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:24:46.625427Z",
     "start_time": "2020-10-24T06:24:46.581435Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.2235)"
      ]
     },
     "execution_count": 117,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = ContrastiveLoss1d(reduction='sum')\n",
    "loss(x1, x2, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:24:47.525104Z",
     "start_time": "2020-10-24T06:24:47.485447Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2.2235],\n",
       "        [0.0000]])"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = ContrastiveLoss1d(reduction='none')\n",
    "loss(x1, x2, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:29:47.060876Z",
     "start_time": "2020-10-24T06:29:46.999371Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class ContrastiveLoss2d(nn.Module):\n",
    "    \n",
    "    @add_docstring(contrastive_loss)\n",
    "    def __init__(self, m=1.25, p=2, reduction='mean'):\n",
    "        \"\"\"OOP version of contrastive loss. We're using the name \"2d\" somewhat\n",
    "        differently here: we use this module to compare 1 image to n different\n",
    "        variants. Picture, for instance, a task where a single example \n",
    "        contains n+1 images and we want to find which of the final n images \n",
    "        are similar to the first image. Concretely, this would be a \n",
    "        multi-label (not multi-class) classification problem with OHE labels.\n",
    "        \n",
    "        The docs for the functional \n",
    "        version are below:\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.m = m\n",
    "        self.p = p\n",
    "        self.loss = partial(contrastive_loss, m=m, p=p, reduction='none')\n",
    "        \n",
    "        if reduction == 'none':\n",
    "            self.reduction = identity\n",
    "        elif reduction == 'row':\n",
    "            self.reduction = partial(torch.sum, dim=-1)\n",
    "        else:\n",
    "            self.reduction = getattr(torch, reduction)\n",
    "        \n",
    "    def forward(self, x1, x2, y_true):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x1: torch.Tensor\n",
    "            Shape (bs, feats).\n",
    "        x2: torch.Tensor\n",
    "            Shape (bs, n_item, n_feats).\n",
    "        y_true:\n",
    "            Shape (bs, n_item).\n",
    "\n",
    "        Returns\n",
    "        ---------\n",
    "        torch.Tensor: scalar if reduction is 'mean' or 'sum', same shape as y\n",
    "        \"\"\"\n",
    "        # if reduction is 'none', or shape (bs,) if reduction is 'row'.\n",
    "        bs, n, dim = x2.shape\n",
    "        res = self.loss(x1.repeat_interleave(n, dim=0), \n",
    "                        x2.view(-1, dim),\n",
    "                        y_true.view(-1, 1))\n",
    "        return self.reduction(res.view(bs, -1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:35:52.558583Z",
     "start_time": "2020-10-24T06:35:52.501135Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.8069, 18.7798, 38.6250, 59.6368, 80.2849],\n",
       "        [24.5225, 25.7855, 25.0192, 25.8827, 23.5852]])"
      ]
     },
     "execution_count": 130,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:37:41.030324Z",
     "start_time": "2020-10-24T06:37:40.966171Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 3, 5])\n",
      "tensor([[[ -0.7902,  18.7978,  38.6418,  59.6638,  80.3098],\n",
      "         [ 15.2447,  60.8061,  63.5002,  85.1861,  91.9863],\n",
      "         [ -0.7529,  18.8695,  38.6729,  59.6541,  80.3533]],\n",
      "\n",
      "        [[ 31.0942,  31.8574,  28.2461,  35.2123,  28.0060],\n",
      "         [ 34.3140,  32.5423,  26.9292,  32.1664,  33.0371],\n",
      "         [531.4734, 534.1118, 531.1953, 532.5379, 529.6300]]])\n"
     ]
    }
   ],
   "source": [
    "# x3 provides 3 examples for each source image in x2. Some are different\n",
    "# and some are similar (notice where noise is amplified vs. dampened).\n",
    "noise = torch.rand(bs, 3, 5) * 10\n",
    "noise[0, [0, -1]] /= 100\n",
    "noise[0, 1] *= 5\n",
    "noise[-1, -1] += 500\n",
    "x3 = x2[:, None, ...] + noise\n",
    "\n",
    "print(x3.shape)\n",
    "print(x3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:37:45.973328Z",
     "start_time": "2020-10-24T06:37:45.918746Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1, 0, 1],\n",
       "        [1, 1, 0]])"
      ]
     },
     "execution_count": 133,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y2d = torch.tensor([[1, 0, 1],\n",
    "                    [1, 1, 0]])\n",
    "y2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:37:48.280563Z",
     "start_time": "2020-10-24T06:37:48.236057Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(39.2557)"
      ]
     },
     "execution_count": 134,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = ContrastiveLoss2d()\n",
    "loss(x2, x3, y2d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:39:07.819890Z",
     "start_time": "2020-10-24T06:39:07.719226Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(235.5344)"
      ]
     },
     "execution_count": 135,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = ContrastiveLoss2d(reduction='sum')\n",
    "loss(x2, x3, y2d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:39:07.819890Z",
     "start_time": "2020-10-24T06:39:07.719226Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1.0237e-02, 2.3552e+02])"
      ]
     },
     "execution_count": 138,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = ContrastiveLoss2d(reduction='row')\n",
    "loss(x2, x3, y2d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-24T06:39:07.819890Z",
     "start_time": "2020-10-24T06:39:07.719226Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.1175e-03, 0.0000e+00, 9.1191e-03],\n",
       "        [9.8526e+01, 1.3700e+02, 0.0000e+00]])"
      ]
     },
     "execution_count": 141,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = ContrastiveLoss2d(reduction='none')\n",
    "loss(x2, x3, y2d)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
