Metric nearness example in Pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n",
      "The line_profiler extension is already loaded. To reload it, use:\n",
      "  %reload_ext line_profiler\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%load_ext line_profiler\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils import data\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "%aimport metrics_pytorch\n",
    "from metrics_pytorch import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's make a simple symmetric graph (3D gridworld)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(125, 375)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph_id = 'simple3d'\n",
    "\n",
    "D = 5\n",
    "N_STATES = D*D*D\n",
    "\n",
    "def idx_to_coord(idx):\n",
    "  return idx // (D*D), (idx % (D*D)) // D, idx % D\n",
    "\n",
    "def coord_to_idx(coord):\n",
    "  x, y, z = coord\n",
    "  return x*D*D + y*D + z\n",
    "\n",
    "def neighbors_c(coord):\n",
    "  x, y, z = coord\n",
    "  return [\n",
    "    ((x + 1) % D, y, z),\n",
    "    ((x - 1) % D, y, z),\n",
    "    (x, (y + 1) % D, z),\n",
    "    (x, (y - 1) % D, z),\n",
    "    (x, y, (z + 1) % D),\n",
    "    (x, y, (z - 1) % D)\n",
    "  ]\n",
    "\n",
    "def neighbors_i(idx):\n",
    "  return list(map(coord_to_idx, neighbors_c(idx_to_coord(idx))))\n",
    "\n",
    "G = nx.Graph()\n",
    "for node in range(N_STATES):\n",
    "  for neighbor in neighbors_i(node):\n",
    "    G.add_edge(node, neighbor, weight=np.random.randint(1, 101))\n",
    "    \n",
    "G = G.to_undirected() # Make symm\n",
    "G.number_of_nodes(), G.number_of_edges()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Since there are only 125 nodes, we will use 1-hot embeddings, and prepare a dataset of $(X_1, X_2, D(X_1,X_2), D(X_1, X_2) + \\text{Noise})$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [],
   "source": [
    "X1, X2, DX, DXnoisy = [], [], [], []\n",
    "\n",
    "for node in range(N_STATES):\n",
    "  for target in range(N_STATES):\n",
    "    X1.append(node)\n",
    "    X2.append(target)\n",
    "    DX.append(nx.astar_path_length(G, node, target))\n",
    "    DXnoisy.append(DX[-1] + np.random.normal(scale=5))\n",
    "\n",
    "X1 = torch.LongTensor(X1).to(DEVICE)\n",
    "X2 = torch.LongTensor(X2).to(DEVICE)\n",
    "DX = torch.FloatTensor(DX).to(DEVICE)\n",
    "DXnoisy = torch.FloatTensor(DXnoisy).to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 251,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LETS USE 20% of the DATA to TRAIN OUR METRICS\n",
    "TRAINING_INDICES = torch.LongTensor(np.random.choice(len(X1), size=(len(X1) // 5,), replace=False)).to(DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now let's define some metrics, and check whether they satisfy triangle inequality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 384,
   "metadata": {},
   "outputs": [],
   "source": [
    "def GroundTruthMetric(x1s, x2s):\n",
    "  idxs = x1s*N_STATES + x2s\n",
    "  return DX[idxs.to(torch.int64)]\n",
    "\n",
    "def NoisyMetric(x1s, x2s):\n",
    "  idxs = x1s*N_STATES + x2s\n",
    "  return DXnoisy[idxs.to(torch.int64)]\n",
    "\n",
    "class MetricModel(nn.Module):\n",
    "  def __init__(self, metric):\n",
    "    super().__init__()\n",
    "    self.metric = metric\n",
    "    self.embedding = nn.Linear(N_STATES, N_STATES, bias=False)\n",
    "    \n",
    "  def forward(self, x, y):\n",
    "    \"\"\"x and y will be ints\"\"\"\n",
    "    x = torch.eye(N_STATES, device=DEVICE)[x]\n",
    "    y = torch.eye(N_STATES, device=DEVICE)[y]\n",
    "    phi_x = self.embedding(x)\n",
    "    phi_y = self.embedding(y)\n",
    "    return self.metric(phi_x, phi_y)\n",
    "\n",
    "mahalanobis  = MetricModel(EuclideanMetric()).to(DEVICE)\n",
    "widenorm     = MetricModel(WideNormMetric(N_STATES, 32, 32)).to(DEVICE)\n",
    "deepnorm     = MetricModel(DeepNormMetric(\n",
    "                  N_STATES, (128, 128), \n",
    "                  activation=lambda: MaxReLUPairwiseActivation(128))).to(DEVICE)\n",
    "neuralmetric = MetricModel(DeepNormMetric(\n",
    "                  N_STATES, (128, 128), activation=lambda: MaxReLUPairwiseActivation(128), \n",
    "                  concave_activation_size=10)).to(DEVICE)\n",
    "notametric   = MetricModel(MLPNonMetric(N_STATES, (128, 128), mode='subtract')).to(DEVICE)\n",
    "\n",
    "def train_metric(metric, opt, epochs=500):\n",
    "  for _ in range(epochs):\n",
    "    loss = F.mse_loss(metric(X1[TRAINING_INDICES], X2[TRAINING_INDICES]), DXnoisy[TRAINING_INDICES])\n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "    \n",
    "def train_on_full_dataset(metric, opt, epochs=500):\n",
    "  for _ in range(epochs):\n",
    "    loss = F.mse_loss(metric(X1, X2), DXnoisy)\n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "      \n",
    "def test_violations(metric, N=10000):\n",
    "  triplets = torch.randint(N_STATES, (N, 3)).to(DEVICE)\n",
    "  start, middle, end = [t.squeeze(1) for t in torch.chunk(triplets, 3, 1)]\n",
    "  violations = metric(start, end) > metric(start, middle) + metric(middle, end) + 1e-4 # buffer for precision\n",
    "  return torch.sum(violations) \n",
    "  \n",
    "def l2_error(metric1, metric2, N=10000):\n",
    "  pairs = torch.randint(N_STATES, (N, 2)).to(DEVICE)\n",
    "  start, end = [t.squeeze(1) for t in torch.chunk(pairs, 2, 1)]\n",
    "  return F.mse_loss(metric1(start, end), metric2(start, end)) \n",
    "\n",
    "def evaluate_metric(metric, N=10000):\n",
    "  print(f'{test_violations(metric, N)} triangle ineq. violations. '+\\\n",
    "        f'{l2_error(metric, GroundTruthMetric, N):.4f} l2 error to GT. '+\\\n",
    "        f'{l2_error(metric, NoisyMetric, N):.4f} l2 error to noisy GT.')\n",
    "  \n",
    "def evaluate_and_train_metric(metric, epochs=20, lr=0.1, training_fn=train_metric):\n",
    "  opt = torch.optim.SGD(metric.parameters(), lr, momentum=0.9, weight_decay=1e-4)\n",
    "  for _ in range(epochs):\n",
    "    evaluate_metric(metric)\n",
    "    training_fn(metric, opt)\n",
    "  evaluate_metric(metric)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 385,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 triangle ineq. violations. 0.0000 l2 error to GT. 24.9282 l2 error to noisy GT.\n"
     ]
    }
   ],
   "source": [
    "evaluate_metric(GroundTruthMetric)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 386,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "306 triangle ineq. violations. 24.3428 l2 error to GT. 0.0000 l2 error to noisy GT.\n"
     ]
    }
   ],
   "source": [
    "evaluate_metric(NoisyMetric)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 387,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 triangle ineq. violations. 10311.9072 l2 error to GT. 10341.7471 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 170.3797 l2 error to GT. 188.9328 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 159.3267 l2 error to GT. 182.2031 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 159.4852 l2 error to GT. 181.7425 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 152.7919 l2 error to GT. 183.4248 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 153.8724 l2 error to GT. 176.5508 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 159.1762 l2 error to GT. 184.9554 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 153.1673 l2 error to GT. 177.6189 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 155.7677 l2 error to GT. 184.5368 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 161.6409 l2 error to GT. 177.2016 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 152.0458 l2 error to GT. 182.0342 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 159.1007 l2 error to GT. 180.2986 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 162.8872 l2 error to GT. 180.4844 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 156.7229 l2 error to GT. 182.4868 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 154.6321 l2 error to GT. 173.9410 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 153.7266 l2 error to GT. 175.6891 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 151.9211 l2 error to GT. 179.8599 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 158.5848 l2 error to GT. 179.3427 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 157.9843 l2 error to GT. 175.7414 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 158.1458 l2 error to GT. 179.6205 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 158.5872 l2 error to GT. 179.4315 l2 error to noisy GT.\n"
     ]
    }
   ],
   "source": [
    "# Performance seems to saturate at around 150-160\n",
    "evaluate_and_train_metric(mahalanobis, lr=1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 388,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 triangle ineq. violations. 10471.7471 l2 error to GT. 10518.0234 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 343.7487 l2 error to GT. 361.9023 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 130.0676 l2 error to GT. 151.8909 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 115.6810 l2 error to GT. 142.0233 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 114.6723 l2 error to GT. 134.4137 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 111.1650 l2 error to GT. 130.0697 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 111.7674 l2 error to GT. 133.4260 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 118.1884 l2 error to GT. 133.4871 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 115.1974 l2 error to GT. 134.8011 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 116.2533 l2 error to GT. 135.4603 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 115.7687 l2 error to GT. 133.2284 l2 error to noisy GT.\n",
      "1 triangle ineq. violations. 115.2697 l2 error to GT. 132.8041 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 112.2026 l2 error to GT. 134.5892 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 112.3192 l2 error to GT. 134.6253 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 114.1099 l2 error to GT. 135.0573 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 116.5683 l2 error to GT. 131.5460 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 113.2129 l2 error to GT. 134.8791 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 111.4643 l2 error to GT. 135.4615 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 113.5648 l2 error to GT. 136.7742 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 112.7072 l2 error to GT. 141.1873 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 114.9576 l2 error to GT. 136.1507 l2 error to noisy GT.\n"
     ]
    }
   ],
   "source": [
    "# Widenorm is much more expressive, with performance saturating around 110-120\n",
    "widenorm     = MetricModel(WideNormMetric(N_STATES, 32, 32)).to(DEVICE)\n",
    "evaluate_and_train_metric(widenorm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 411,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 triangle ineq. violations. 10400.1523 l2 error to GT. 10573.5264 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 168.0085 l2 error to GT. 193.4304 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 138.3754 l2 error to GT. 162.1016 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 131.3053 l2 error to GT. 158.0584 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 131.4579 l2 error to GT. 156.5459 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 127.7473 l2 error to GT. 149.6919 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 128.4855 l2 error to GT. 144.8129 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 122.9368 l2 error to GT. 144.8664 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 121.3388 l2 error to GT. 140.4595 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 125.2140 l2 error to GT. 135.5015 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 124.6440 l2 error to GT. 145.1458 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 120.3016 l2 error to GT. 145.5744 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 123.2092 l2 error to GT. 142.5954 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 121.3615 l2 error to GT. 140.7750 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 122.8521 l2 error to GT. 138.3803 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 121.5707 l2 error to GT. 139.4923 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 120.5993 l2 error to GT. 143.8555 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 119.5183 l2 error to GT. 142.4641 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 120.2152 l2 error to GT. 138.1774 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 117.2440 l2 error to GT. 137.6893 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 118.6023 l2 error to GT. 145.5281 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 119.5503 l2 error to GT. 138.4540 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 116.7727 l2 error to GT. 139.7199 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 115.6949 l2 error to GT. 136.4657 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 120.0111 l2 error to GT. 137.0289 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 114.7313 l2 error to GT. 140.2489 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 119.4570 l2 error to GT. 133.0388 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 118.3130 l2 error to GT. 134.4594 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 113.1618 l2 error to GT. 141.5204 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 116.1219 l2 error to GT. 132.6297 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 114.3754 l2 error to GT. 138.9150 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 117.5036 l2 error to GT. 135.0357 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 118.2232 l2 error to GT. 140.0848 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 113.8021 l2 error to GT. 137.4796 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 121.7386 l2 error to GT. 135.7046 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 118.7077 l2 error to GT. 142.5493 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 116.2255 l2 error to GT. 135.5312 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 120.9782 l2 error to GT. 137.7956 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 116.5809 l2 error to GT. 138.2231 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 117.9624 l2 error to GT. 139.5633 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 113.5770 l2 error to GT. 137.8150 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 118.2524 l2 error to GT. 137.2758 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 120.4761 l2 error to GT. 141.1592 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 121.4365 l2 error to GT. 141.5090 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 118.1553 l2 error to GT. 136.4307 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 112.8365 l2 error to GT. 141.3717 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 115.8521 l2 error to GT. 135.6176 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 123.5960 l2 error to GT. 136.2330 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 117.2447 l2 error to GT. 134.5593 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 114.0045 l2 error to GT. 136.0889 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 116.7445 l2 error to GT. 134.9068 l2 error to noisy GT.\n"
     ]
    }
   ],
   "source": [
    "# Requires a smaller learning rate and trains a bit slower. \n",
    "deepnorm     = MetricModel(DeepNormMetric(\n",
    "                  N_STATES, (128, 128), \n",
    "                  activation=lambda: MaxReLUPairwiseActivation(128))).to(DEVICE)\n",
    "evaluate_and_train_metric(deepnorm, epochs=50, lr=1e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 412,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 triangle ineq. violations. 10474.0420 l2 error to GT. 10534.6318 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 164.4842 l2 error to GT. 186.1844 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 135.1067 l2 error to GT. 154.1147 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 135.4845 l2 error to GT. 150.5378 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 119.1945 l2 error to GT. 145.5924 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 123.3479 l2 error to GT. 142.1933 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 124.8420 l2 error to GT. 143.2315 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 121.7935 l2 error to GT. 139.9626 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 115.8865 l2 error to GT. 142.2808 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 117.8065 l2 error to GT. 144.5377 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 118.9322 l2 error to GT. 137.1372 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 122.7727 l2 error to GT. 140.6126 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 121.9829 l2 error to GT. 136.0816 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 119.6948 l2 error to GT. 137.9465 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 121.8991 l2 error to GT. 138.3591 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 120.7866 l2 error to GT. 133.9144 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 122.5702 l2 error to GT. 139.2288 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 118.6030 l2 error to GT. 142.9029 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 117.6856 l2 error to GT. 139.4162 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 120.6760 l2 error to GT. 143.7584 l2 error to noisy GT.\n",
      "0 triangle ineq. violations. 122.4888 l2 error to GT. 144.8813 l2 error to noisy GT.\n"
     ]
    }
   ],
   "source": [
    "# Concave activations might help a bit here. \n",
    "neuralmetric = MetricModel(DeepNormMetric(\n",
    "                  N_STATES, (128, 128), activation=lambda: MaxReLUPairwiseActivation(128), \n",
    "                  concave_activation_size=10)).to(DEVICE)\n",
    "evaluate_and_train_metric(neuralmetric, epochs=20, lr=1e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 413,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10000 triangle ineq. violations. 10495.0625 l2 error to GT. 10534.7617 l2 error to noisy GT.\n",
      "77 triangle ineq. violations. 184.6394 l2 error to GT. 210.0858 l2 error to noisy GT.\n",
      "148 triangle ineq. violations. 178.7564 l2 error to GT. 193.9753 l2 error to noisy GT.\n",
      "187 triangle ineq. violations. 181.3610 l2 error to GT. 198.7089 l2 error to noisy GT.\n",
      "196 triangle ineq. violations. 193.4611 l2 error to GT. 202.2033 l2 error to noisy GT.\n",
      "189 triangle ineq. violations. 185.0388 l2 error to GT. 208.8876 l2 error to noisy GT.\n",
      "194 triangle ineq. violations. 183.0283 l2 error to GT. 209.8664 l2 error to noisy GT.\n",
      "173 triangle ineq. violations. 195.1491 l2 error to GT. 204.6775 l2 error to noisy GT.\n",
      "203 triangle ineq. violations. 185.0747 l2 error to GT. 207.7710 l2 error to noisy GT.\n",
      "191 triangle ineq. violations. 194.9169 l2 error to GT. 208.4712 l2 error to noisy GT.\n",
      "235 triangle ineq. violations. 187.6752 l2 error to GT. 203.4320 l2 error to noisy GT.\n",
      "194 triangle ineq. violations. 186.5423 l2 error to GT. 198.7361 l2 error to noisy GT.\n",
      "219 triangle ineq. violations. 187.4881 l2 error to GT. 207.1767 l2 error to noisy GT.\n",
      "187 triangle ineq. violations. 199.9080 l2 error to GT. 209.6958 l2 error to noisy GT.\n",
      "189 triangle ineq. violations. 191.7513 l2 error to GT. 209.5724 l2 error to noisy GT.\n",
      "222 triangle ineq. violations. 192.1108 l2 error to GT. 207.9411 l2 error to noisy GT.\n",
      "194 triangle ineq. violations. 193.4354 l2 error to GT. 204.9565 l2 error to noisy GT.\n",
      "222 triangle ineq. violations. 194.8017 l2 error to GT. 207.4061 l2 error to noisy GT.\n",
      "211 triangle ineq. violations. 189.6991 l2 error to GT. 204.2804 l2 error to noisy GT.\n",
      "196 triangle ineq. violations. 192.1049 l2 error to GT. 200.2528 l2 error to noisy GT.\n",
      "212 triangle ineq. violations. 192.7870 l2 error to GT. 204.3376 l2 error to noisy GT.\n"
     ]
    }
   ],
   "source": [
    "# MLP needs a very small learning rate to train. \n",
    "# However, without the inductive bias, it doesn't fit very well.\n",
    "# And it violates triangle inequality\n",
    "notametric   = MetricModel(MLPNonMetric(N_STATES, (128, 128), mode='subtract')).to(DEVICE)\n",
    "evaluate_and_train_metric(notametric, lr=3e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 414,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 triangle ineq. violations. 10539.4316 l2 error to GT. 10582.9150 l2 error to noisy GT.\n",
      "101 triangle ineq. violations. 83.6959 l2 error to GT. 102.5730 l2 error to noisy GT.\n",
      "118 triangle ineq. violations. 30.7381 l2 error to GT. 47.8616 l2 error to noisy GT.\n",
      "131 triangle ineq. violations. 23.0263 l2 error to GT. 36.3405 l2 error to noisy GT.\n",
      "173 triangle ineq. violations. 17.7455 l2 error to GT. 27.6664 l2 error to noisy GT.\n",
      "154 triangle ineq. violations. 14.4404 l2 error to GT. 22.5303 l2 error to noisy GT.\n",
      "154 triangle ineq. violations. 20.2875 l2 error to GT. 24.9125 l2 error to noisy GT.\n",
      "167 triangle ineq. violations. 13.2997 l2 error to GT. 16.3037 l2 error to noisy GT.\n",
      "184 triangle ineq. violations. 29.3111 l2 error to GT. 30.5212 l2 error to noisy GT.\n",
      "191 triangle ineq. violations. 16.2492 l2 error to GT. 14.1748 l2 error to noisy GT.\n",
      "243 triangle ineq. violations. 24.8715 l2 error to GT. 20.8237 l2 error to noisy GT.\n",
      "179 triangle ineq. violations. 19.3138 l2 error to GT. 13.8690 l2 error to noisy GT.\n",
      "201 triangle ineq. violations. 16.7088 l2 error to GT. 9.5283 l2 error to noisy GT.\n",
      "213 triangle ineq. violations. 16.9525 l2 error to GT. 8.6766 l2 error to noisy GT.\n",
      "177 triangle ineq. violations. 19.3179 l2 error to GT. 8.8804 l2 error to noisy GT.\n",
      "226 triangle ineq. violations. 17.9997 l2 error to GT. 6.2329 l2 error to noisy GT.\n",
      "217 triangle ineq. violations. 22.4090 l2 error to GT. 10.1163 l2 error to noisy GT.\n",
      "169 triangle ineq. violations. 19.8768 l2 error to GT. 6.6022 l2 error to noisy GT.\n",
      "221 triangle ineq. violations. 19.7148 l2 error to GT. 6.3026 l2 error to noisy GT.\n",
      "190 triangle ineq. violations. 19.3832 l2 error to GT. 5.0075 l2 error to noisy GT.\n",
      "214 triangle ineq. violations. 20.0980 l2 error to GT. 4.8041 l2 error to noisy GT.\n"
     ]
    }
   ],
   "source": [
    "# Note that if we let it train on the whole dataset, it fits exceptionally well (though has violations)\n",
    "notametric   = MetricModel(MLPNonMetric(N_STATES, (128, 128), mode='subtract')).to(DEVICE)\n",
    "evaluate_and_train_metric(notametric, lr=3e-4, training_fn=train_on_full_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}