test_tsp_transformer_beamsearch_TSP50.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Transformer Network for the Traveling Salesman Problem\n",
    "\n",
    "Xavier Bresson, Thomas Laurent, Feb 2021<br>\n",
    "\n",
    "Arxiv : https://arxiv.org/pdf/2103.03012.pdf<br>\n",
    "Talk : https://ipam.wistia.com/medias/0jrweluovs<br>\n",
    "Slides : https://t.co/ySxGiKtQL5<br>\n",
    "\n",
    "This code decodes TSP solutions with beam search given a trained network.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "###################\n",
    "# Libs\n",
    "###################\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import time\n",
    "import argparse\n",
    "\n",
    "import os\n",
    "import datetime\n",
    "\n",
    "from torch.distributions.categorical import Categorical\n",
    "\n",
    "# visualization \n",
    "%matplotlib inline\n",
    "from IPython.display import set_matplotlib_formats, clear_output\n",
    "set_matplotlib_formats('png2x','pdf')\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "try: \n",
    "    import networkx as nx\n",
    "    from scipy.spatial.distance import pdist, squareform\n",
    "    from concorde.tsp import TSPSolver # !pip install -e pyconcorde\n",
    "except:\n",
    "    pass\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "###################\n",
    "# Hardware : CPU / GPU(s)\n",
    "###################\n",
    "\n",
    "device = torch.device(\"cpu\"); gpu_id = -1 # select CPU\n",
    "\n",
    "gpu_id = '0' # select a single GPU  \n",
    "#gpu_id = '0,1' # select multiple GPUs  \n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)  \n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    print('GPU name: {:s}, gpu_id: {:s}'.format(torch.cuda.get_device_name(0),gpu_id))   \n",
    "    \n",
    "print(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "###################\n",
    "# Hyper-parameters\n",
    "###################\n",
    "\n",
    "class DotDict(dict):\n",
    "    def __init__(self, **kwds):\n",
    "        self.update(kwds)\n",
    "        self.__dict__ = self\n",
    "        \n",
    "args = DotDict()\n",
    "args.nb_nodes = 20 # TSP20\n",
    "args.nb_nodes = 50 # TSP50\n",
    "# args.nb_nodes = 100 # TSP100\n",
    "args.bsz = 512 # TSP20 TSP50\n",
    "args.dim_emb = 128\n",
    "args.dim_ff = 512\n",
    "args.dim_input_nodes = 2\n",
    "args.nb_layers_encoder = 6\n",
    "args.nb_layers_decoder = 2\n",
    "args.nb_heads = 8\n",
    "args.nb_epochs = 10000\n",
    "args.nb_batch_per_epoch = 2500\n",
    "args.nb_batch_eval = 20\n",
    "args.gpu_id = gpu_id\n",
    "args.lr = 1e-4\n",
    "args.tol = 1e-3\n",
    "args.batchnorm = True  # if batchnorm=True  than batch norm is used\n",
    "#args.batchnorm = False # if batchnorm=False than layer norm is used\n",
    "args.max_len_PE = 1000\n",
    "\n",
    "print(args)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "###################\n",
    "# Network definition\n",
    "# Notation : \n",
    "#            bsz : batch size\n",
    "#            nb_nodes : number of nodes/cities\n",
    "#            dim_emb : embedding/hidden dimension\n",
    "#            nb_heads : nb of attention heads\n",
    "#            dim_ff : feed-forward dimension\n",
    "#            nb_layers : number of encoder/decoder layers\n",
    "###################\n",
    "def compute_tour_length(x, tour): \n",
    "    \"\"\"\n",
    "    Compute the length of a batch of tours\n",
    "    Inputs : x of size (bsz, nb_nodes, 2) batch of tsp tour instances\n",
    "             tour of size (bsz, nb_nodes) batch of sequences (node indices) of tsp tours\n",
    "    Output : L of size (bsz,)             batch of lengths of each tsp tour\n",
    "    \"\"\"\n",
    "    bsz = x.shape[0]\n",
    "    nb_nodes = x.shape[1]\n",
    "    arange_vec = torch.arange(bsz, device=x.device)\n",
    "    first_cities = x[arange_vec, tour[:,0], :] # size(first_cities)=(bsz,2)\n",
    "    previous_cities = first_cities\n",
    "    L = torch.zeros(bsz, device=x.device)\n",
    "    with torch.no_grad():\n",
    "        for i in range(1,nb_nodes):\n",
    "            current_cities = x[arange_vec, tour[:,i], :] \n",
    "            L += torch.sum( (current_cities - previous_cities)**2 , dim=1 )**0.5 # dist(current, previous node) \n",
    "            previous_cities = current_cities\n",
    "        L += torch.sum( (current_cities - first_cities)**2 , dim=1 )**0.5 # dist(last, first node)  \n",
    "    return L\n",
    "\n",
    "\n",
    "class Transformer_encoder_net(nn.Module):\n",
    "    \"\"\"\n",
    "    Encoder network based on self-attention transformer\n",
    "    Inputs :  \n",
    "      h of size      (bsz, nb_nodes+1, dim_emb)    batch of input cities\n",
    "    Outputs :  \n",
    "      h of size      (bsz, nb_nodes+1, dim_emb)    batch of encoded cities\n",
    "      score of size  (bsz, nb_nodes+1, nb_nodes+1) batch of attention scores\n",
    "    \"\"\"\n",
    "    def __init__(self, nb_layers, dim_emb, nb_heads, dim_ff, batchnorm):\n",
    "        super(Transformer_encoder_net, self).__init__()\n",
    "        assert dim_emb == nb_heads* (dim_emb//nb_heads) # check if dim_emb is divisible by nb_heads\n",
    "        self.MHA_layers = nn.ModuleList( [nn.MultiheadAttention(dim_emb, nb_heads) for _ in range(nb_layers)] )\n",
    "        self.linear1_layers = nn.ModuleList( [nn.Linear(dim_emb, dim_ff) for _ in range(nb_layers)] )\n",
    "        self.linear2_layers = nn.ModuleList( [nn.Linear(dim_ff, dim_emb) for _ in range(nb_layers)] )   \n",
    "        if batchnorm:\n",
    "            self.norm1_layers = nn.ModuleList( [nn.BatchNorm1d(dim_emb) for _ in range(nb_layers)] )\n",
    "            self.norm2_layers = nn.ModuleList( [nn.BatchNorm1d(dim_emb) for _ in range(nb_layers)] )\n",
    "        else:\n",
    "            self.norm1_layers = nn.ModuleList( [nn.LayerNorm(dim_emb) for _ in range(nb_layers)] )\n",
    "            self.norm2_layers = nn.ModuleList( [nn.LayerNorm(dim_emb) for _ in range(nb_layers)] )\n",
    "        self.nb_layers = nb_layers\n",
    "        self.nb_heads = nb_heads\n",
    "        self.batchnorm = batchnorm\n",
    "        \n",
    "    def forward(self, h):      \n",
    "        # PyTorch nn.MultiheadAttention requires input size (seq_len, bsz, dim_emb) \n",
    "        h = h.transpose(0,1) # size(h)=(nb_nodes, bsz, dim_emb)  \n",
    "        # L layers\n",
    "        for i in range(self.nb_layers):\n",
    "            h_rc = h # residual connection, size(h_rc)=(nb_nodes, bsz, dim_emb)\n",
    "            h, score = self.MHA_layers[i](h, h, h) # size(h)=(nb_nodes, bsz, dim_emb), size(score)=(bsz, nb_nodes, nb_nodes)\n",
    "            # add residual connection\n",
    "            h = h_rc + h # size(h)=(nb_nodes, bsz, dim_emb)\n",
    "            if self.batchnorm:\n",
    "                # Pytorch nn.BatchNorm1d requires input size (bsz, dim, seq_len)\n",
    "                h = h.permute(1,2,0).contiguous() # size(h)=(bsz, dim_emb, nb_nodes)\n",
    "                h = self.norm1_layers[i](h)       # size(h)=(bsz, dim_emb, nb_nodes)\n",
    "                h = h.permute(2,0,1).contiguous() # size(h)=(nb_nodes, bsz, dim_emb)\n",
    "            else:\n",
    "                h = self.norm1_layers[i](h)       # size(h)=(nb_nodes, bsz, dim_emb) \n",
    "            # feedforward\n",
    "            h_rc = h # residual connection\n",
    "            h = self.linear2_layers[i](torch.relu(self.linear1_layers[i](h)))\n",
    "            h = h_rc + h # size(h)=(nb_nodes, bsz, dim_emb)\n",
    "            if self.batchnorm:\n",
    "                h = h.permute(1,2,0).contiguous() # size(h)=(bsz, dim_emb, nb_nodes)\n",
    "                h = self.norm2_layers[i](h)       # size(h)=(bsz, dim_emb, nb_nodes)\n",
    "                h = h.permute(2,0,1).contiguous() # size(h)=(nb_nodes, bsz, dim_emb)\n",
    "            else:\n",
    "                h = self.norm2_layers[i](h) # size(h)=(nb_nodes, bsz, dim_emb)\n",
    "        # Transpose h\n",
    "        h = h.transpose(0,1) # size(h)=(bsz, nb_nodes, dim_emb)\n",
    "        return h, score\n",
    "    \n",
    "\n",
    "def myMHA(Q, K, V, nb_heads, mask=None, clip_value=None):\n",
    "    \"\"\"\n",
    "    Compute multi-head attention (MHA) given a query Q, key K, value V and attention mask :\n",
    "      h = Concat_{k=1}^nb_heads softmax(Q_k^T.K_k).V_k \n",
    "    Note : We did not use nn.MultiheadAttention to avoid re-computing all linear transformations at each call.\n",
    "    Inputs : Q of size (bsz, dim_emb, 1)                batch of queries\n",
    "             K of size (bsz, dim_emb, nb_nodes+1)       batch of keys\n",
    "             V of size (bsz, dim_emb, nb_nodes+1)       batch of values\n",
    "             mask of size (bsz, nb_nodes+1)             batch of masks of visited cities\n",
    "             clip_value is a scalar \n",
    "    Outputs : attn_output of size (bsz, 1, dim_emb)     batch of attention vectors\n",
    "              attn_weights of size (bsz, 1, nb_nodes+1) batch of attention weights\n",
    "    \"\"\"\n",
    "    bsz, nb_nodes, emd_dim = K.size() #  dim_emb must be divisable by nb_heads\n",
    "    if nb_heads>1:\n",
    "        # PyTorch view requires contiguous dimensions for correct reshaping\n",
    "        Q = Q.transpose(1,2).contiguous() # size(Q)=(bsz, dim_emb, 1)\n",
    "        Q = Q.view(bsz*nb_heads, emd_dim//nb_heads, 1) # size(Q)=(bsz*nb_heads, dim_emb//nb_heads, 1)\n",
    "        Q = Q.transpose(1,2).contiguous() # size(Q)=(bsz*nb_heads, 1, dim_emb//nb_heads)\n",
    "        K = K.transpose(1,2).contiguous() # size(K)=(bsz, dim_emb, nb_nodes+1)\n",
    "        K = K.view(bsz*nb_heads, emd_dim//nb_heads, nb_nodes) # size(K)=(bsz*nb_heads, dim_emb//nb_heads, nb_nodes+1)\n",
    "        K = K.transpose(1,2).contiguous() # size(K)=(bsz*nb_heads, nb_nodes+1, dim_emb//nb_heads)\n",
    "        V = V.transpose(1,2).contiguous() # size(V)=(bsz, dim_emb, nb_nodes+1)\n",
    "        V = V.view(bsz*nb_heads, emd_dim//nb_heads, nb_nodes) # size(V)=(bsz*nb_heads, dim_emb//nb_heads, nb_nodes+1)\n",
    "        V = V.transpose(1,2).contiguous() # size(V)=(bsz*nb_heads, nb_nodes+1, dim_emb//nb_heads)\n",
    "    attn_weights = torch.bmm(Q, K.transpose(1,2))/ Q.size(-1)**0.5 # size(attn_weights)=(bsz*nb_heads, 1, nb_nodes+1)\n",
    "    if clip_value is not None:\n",
    "        attn_weights = clip_value * torch.tanh(attn_weights)\n",
    "    if mask is not None:\n",
    "        if nb_heads>1:\n",
    "            mask = torch.repeat_interleave(mask, repeats=nb_heads, dim=0) # size(mask)=(bsz*nb_heads, nb_nodes+1)\n",
    "        #attn_weights = attn_weights.masked_fill(mask.unsqueeze(1), float('-inf')) # size(attn_weights)=(bsz*nb_heads, 1, nb_nodes+1)\n",
    "        attn_weights = attn_weights.masked_fill(mask.unsqueeze(1), float('-1e9')) # size(attn_weights)=(bsz*nb_heads, 1, nb_nodes+1)\n",
    "    attn_weights = torch.softmax(attn_weights, dim=-1) # size(attn_weights)=(bsz*nb_heads, 1, nb_nodes+1)\n",
    "    attn_output = torch.bmm(attn_weights, V) # size(attn_output)=(bsz*nb_heads, 1, dim_emb//nb_heads)\n",
    "    if nb_heads>1:\n",
    "        attn_output = attn_output.transpose(1,2).contiguous() # size(attn_output)=(bsz*nb_heads, dim_emb//nb_heads, 1)\n",
    "        attn_output = attn_output.view(bsz, emd_dim, 1) # size(attn_output)=(bsz, dim_emb, 1)\n",
    "        attn_output = attn_output.transpose(1,2).contiguous() # size(attn_output)=(bsz, 1, dim_emb)\n",
    "        attn_weights = attn_weights.view(bsz, nb_heads, 1, nb_nodes) # size(attn_weights)=(bsz, nb_heads, 1, nb_nodes+1)\n",
    "        attn_weights = attn_weights.mean(dim=1) # mean over the heads, size(attn_weights)=(bsz, 1, nb_nodes+1)\n",
    "    return attn_output, attn_weights\n",
    "    \n",
    "    \n",
    "class AutoRegressiveDecoderLayer(nn.Module):\n",
    "    \"\"\"\n",
    "    Single decoder layer based on self-attention and query-attention\n",
    "    Inputs :  \n",
    "      h_t of size      (bsz, 1, dim_emb)          batch of input queries\n",
    "      K_att of size    (bsz, nb_nodes+1, dim_emb) batch of query-attention keys\n",
    "      V_att of size    (bsz, nb_nodes+1, dim_emb) batch of query-attention values\n",
    "      mask of size     (bsz, nb_nodes+1)          batch of masks of visited cities\n",
    "    Output :  \n",
    "      h_t of size (bsz, nb_nodes+1)               batch of transformed queries\n",
    "    \"\"\"\n",
    "    def __init__(self, dim_emb, nb_heads):\n",
    "        super(AutoRegressiveDecoderLayer, self).__init__()\n",
    "        self.dim_emb = dim_emb\n",
    "        self.nb_heads = nb_heads\n",
    "        self.Wq_selfatt = nn.Linear(dim_emb, dim_emb)\n",
    "        self.Wk_selfatt = nn.Linear(dim_emb, dim_emb)\n",
    "        self.Wv_selfatt = nn.Linear(dim_emb, dim_emb)\n",
    "        self.W0_selfatt = nn.Linear(dim_emb, dim_emb)\n",
    "        self.W0_att = nn.Linear(dim_emb, dim_emb)\n",
    "        self.Wq_att = nn.Linear(dim_emb, dim_emb)\n",
    "        self.W1_MLP = nn.Linear(dim_emb, dim_emb)\n",
    "        self.W2_MLP = nn.Linear(dim_emb, dim_emb)\n",
    "        self.BN_selfatt = nn.LayerNorm(dim_emb)\n",
    "        self.BN_att = nn.LayerNorm(dim_emb)\n",
    "        self.BN_MLP = nn.LayerNorm(dim_emb)\n",
    "        self.K_sa = None\n",
    "        self.V_sa = None\n",
    "\n",
    "    def reset_selfatt_keys_values(self):\n",
    "        self.K_sa = None\n",
    "        self.V_sa = None\n",
    "        \n",
    "    # For beam search\n",
    "    def reorder_selfatt_keys_values(self, t, idx_top_beams):\n",
    "        bsz, B = idx_top_beams.size()\n",
    "        zero_to_B = torch.arange(B, device=idx_top_beams.device) # [0,1,...,B-1]\n",
    "        B2 = self.K_sa.size(0)// bsz\n",
    "        self.K_sa = self.K_sa.view(bsz, B2, t+1, self.dim_emb) # size(self.K_sa)=(bsz, B2, t+1, dim_emb)\n",
    "        K_sa_tmp = self.K_sa.clone()\n",
    "        self.K_sa = torch.zeros(bsz, B, t+1, self.dim_emb, device=idx_top_beams.device)\n",
    "        for b in range(bsz):\n",
    "            self.K_sa[b, zero_to_B, :, :] = K_sa_tmp[b, idx_top_beams[b], :, :]\n",
    "        self.K_sa = self.K_sa.view(bsz*B, t+1, self.dim_emb) # size(self.K_sa)=(bsz*B, t+1, dim_emb)\n",
    "        self.V_sa = self.V_sa.view(bsz, B2, t+1, self.dim_emb) # size(self.K_sa)=(bsz, B, t+1, dim_emb)\n",
    "        V_sa_tmp = self.V_sa.clone()\n",
    "        self.V_sa = torch.zeros(bsz, B, t+1, self.dim_emb, device=idx_top_beams.device)\n",
    "        for b in range(bsz):\n",
    "            self.V_sa[b, zero_to_B, :, :] = V_sa_tmp[b, idx_top_beams[b], :, :]\n",
    "        self.V_sa = self.V_sa.view(bsz*B, t+1, self.dim_emb) # size(self.K_sa)=(bsz*B, t+1, dim_emb)\n",
    "\n",
    "    # For beam search\n",
    "    def repeat_selfatt_keys_values(self, B):\n",
    "        self.K_sa = torch.repeat_interleave(self.K_sa, B, dim=0) # size(self.K_sa)=(bsz.B, t+1, dim_emb)\n",
    "        self.V_sa = torch.repeat_interleave(self.V_sa, B, dim=0) # size(self.K_sa)=(bsz.B, t+1, dim_emb)\n",
    "        \n",
    "    def forward(self, h_t, K_att, V_att, mask):\n",
    "        bsz = h_t.size(0)\n",
    "        h_t = h_t.view(bsz,1,self.dim_emb) # size(h_t)=(bsz, 1, dim_emb)\n",
    "        # embed the query for self-attention\n",
    "        q_sa = self.Wq_selfatt(h_t) # size(q_sa)=(bsz, 1, dim_emb)\n",
    "        k_sa = self.Wk_selfatt(h_t) # size(k_sa)=(bsz, 1, dim_emb)\n",
    "        v_sa = self.Wv_selfatt(h_t) # size(v_sa)=(bsz, 1, dim_emb)\n",
    "        # concatenate the new self-attention key and value to the previous keys and values\n",
    "        if self.K_sa is None:\n",
    "            self.K_sa = k_sa # size(self.K_sa)=(bsz, 1, dim_emb)\n",
    "            self.V_sa = v_sa # size(self.V_sa)=(bsz, 1, dim_emb)\n",
    "        else:\n",
    "            self.K_sa = torch.cat([self.K_sa, k_sa], dim=1)\n",
    "            self.V_sa = torch.cat([self.V_sa, v_sa], dim=1)\n",
    "        # compute self-attention between nodes in the partial tour\n",
    "        h_t = h_t + self.W0_selfatt( myMHA(q_sa, self.K_sa, self.V_sa, self.nb_heads)[0] ) # size(h_t)=(bsz, 1, dim_emb)\n",
    "        h_t = self.BN_selfatt(h_t.squeeze()) # size(h_t)=(bsz, dim_emb)\n",
    "        h_t = h_t.view(bsz, 1, self.dim_emb) # size(h_t)=(bsz, 1, dim_emb)\n",
    "        # compute attention between self-attention nodes and encoding nodes in the partial tour (translation process)\n",
    "        q_a = self.Wq_att(h_t) # size(q_a)=(bsz, 1, dim_emb)\n",
    "        h_t = h_t + self.W0_att( myMHA(q_a, K_att, V_att, self.nb_heads, mask)[0] ) # size(h_t)=(bsz, 1, dim_emb)\n",
    "        h_t = self.BN_att(h_t.squeeze()) # size(h_t)=(bsz, dim_emb)\n",
    "        h_t = h_t.view(bsz, 1, self.dim_emb) # size(h_t)=(bsz, 1, dim_emb)\n",
    "        # MLP\n",
    "        h_t = h_t + self.W2_MLP(torch.relu(self.W1_MLP(h_t)))\n",
    "        h_t = self.BN_MLP(h_t.squeeze(1)) # size(h_t)=(bsz, dim_emb)\n",
    "        return h_t\n",
    "        \n",
    "        \n",
    "class Transformer_decoder_net(nn.Module): \n",
    "    \"\"\"\n",
    "    Decoder network based on self-attention and query-attention transformers\n",
    "    Inputs :  \n",
    "      h_t of size      (bsz, 1, dim_emb)                            batch of input queries\n",
    "      K_att of size    (bsz, nb_nodes+1, dim_emb*nb_layers_decoder) batch of query-attention keys for all decoding layers\n",
    "      V_att of size    (bsz, nb_nodes+1, dim_emb*nb_layers_decoder) batch of query-attention values for all decoding layers\n",
    "      mask of size     (bsz, nb_nodes+1)                            batch of masks of visited cities\n",
    "    Output :  \n",
    "      prob_next_node of size (bsz, nb_nodes+1)                      batch of probabilities of next node\n",
    "    \"\"\"\n",
    "    def __init__(self, dim_emb, nb_heads, nb_layers_decoder):\n",
    "        super(Transformer_decoder_net, self).__init__()\n",
    "        self.dim_emb = dim_emb\n",
    "        self.nb_heads = nb_heads\n",
    "        self.nb_layers_decoder = nb_layers_decoder\n",
    "        self.decoder_layers = nn.ModuleList( [AutoRegressiveDecoderLayer(dim_emb, nb_heads) for _ in range(nb_layers_decoder-1)] )\n",
    "        self.Wq_final = nn.Linear(dim_emb, dim_emb)\n",
    "        \n",
    "    # Reset to None self-attention keys and values when decoding starts \n",
    "    def reset_selfatt_keys_values(self): \n",
    "        for l in range(self.nb_layers_decoder-1):\n",
    "            self.decoder_layers[l].reset_selfatt_keys_values()\n",
    "            \n",
    "    # For beam search\n",
    "    def reorder_selfatt_keys_values(self, t, idx_top_beams):\n",
    "        for l in range(self.nb_layers_decoder-1):\n",
    "            self.decoder_layers[l].reorder_selfatt_keys_values(t, idx_top_beams)\n",
    "    \n",
    "    # For beam search\n",
    "    def repeat_selfatt_keys_values(self, B):\n",
    "        for l in range(self.nb_layers_decoder-1):\n",
    "            self.decoder_layers[l].repeat_selfatt_keys_values(B)\n",
    "     \n",
    "    def forward(self, h_t, K_att, V_att, mask):\n",
    "        for l in range(self.nb_layers_decoder):\n",
    "            K_att_l = K_att[:,:,l*self.dim_emb:(l+1)*self.dim_emb].contiguous()  # size(K_att_l)=(bsz, nb_nodes+1, dim_emb)\n",
    "            V_att_l = V_att[:,:,l*self.dim_emb:(l+1)*self.dim_emb].contiguous()  # size(V_att_l)=(bsz, nb_nodes+1, dim_emb)\n",
    "            if l<self.nb_layers_decoder-1: # decoder layers with multiple heads (intermediate layers)\n",
    "                h_t = self.decoder_layers[l](h_t, K_att_l, V_att_l, mask)\n",
    "            else: # decoder layers with single head (final layer)\n",
    "                q_final = self.Wq_final(h_t)\n",
    "                bsz = h_t.size(0)\n",
    "                q_final = q_final.view(bsz, 1, self.dim_emb)\n",
    "                attn_weights = myMHA(q_final, K_att_l, V_att_l, 1, mask, 10)[1] \n",
    "        prob_next_node = attn_weights.squeeze(1) \n",
    "        return prob_next_node\n",
    "\n",
    "\n",
    "def generate_positional_encoding(d_model, max_len):\n",
    "    \"\"\"\n",
    "    Create standard transformer PEs.\n",
    "    Inputs :  \n",
    "      d_model is a scalar correspoding to the hidden dimension\n",
    "      max_len is the maximum length of the sequence\n",
    "    Output :  \n",
    "      pe of size (max_len, d_model), where d_model=dim_emb, max_len=1000\n",
    "    \"\"\"\n",
    "    pe = torch.zeros(max_len, d_model)\n",
    "    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
    "    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))\n",
    "    pe[:,0::2] = torch.sin(position * div_term)\n",
    "    pe[:,1::2] = torch.cos(position * div_term)\n",
    "    return pe\n",
    "    \n",
    "    \n",
    "class TSP_net(nn.Module): \n",
    "    \"\"\"\n",
    "    The TSP network is composed of two steps :\n",
    "      Step 1. Encoder step : Take a set of 2D points representing a fully connected graph \n",
    "                             and encode the set with self-transformer.\n",
    "      Step 2. Decoder step : Build the TSP tour recursively/autoregressively, \n",
    "                             i.e. one node at a time, with a self-transformer and query-transformer. \n",
    "    Inputs : \n",
    "      x of size (bsz, nb_nodes, dim_emb) Euclidian coordinates of the nodes/cities\n",
    "      deterministic is a boolean : If True the salesman will chose the city with highest probability. \n",
    "                                   If False the salesman will chose the city with Bernouilli sampling.\n",
    "    Outputs : \n",
    "      tours of size (bsz, nb_nodes) : batch of tours, i.e. sequences of ordered cities \n",
    "                                      tours[b,t] contains the idx of the city visited at step t in batch b\n",
    "      sumLogProbOfActions of size (bsz,) : batch of sum_t log prob( pi_t | pi_(t-1),...,pi_0 )\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, dim_input_nodes, dim_emb, dim_ff, nb_layers_encoder, nb_layers_decoder, nb_heads, max_len_PE,\n",
    "                 batchnorm=True):\n",
    "        super(TSP_net, self).__init__()\n",
    "        \n",
    "        self.dim_emb = dim_emb\n",
    "        \n",
    "        # input embedding layer\n",
    "        self.input_emb = nn.Linear(dim_input_nodes, dim_emb)\n",
    "        \n",
    "        # encoder layer\n",
    "        self.encoder = Transformer_encoder_net(nb_layers_encoder, dim_emb, nb_heads, dim_ff, batchnorm)\n",
    "        \n",
    "        # vector to start decoding \n",
    "        self.start_placeholder = nn.Parameter(torch.randn(dim_emb))\n",
    "        \n",
    "        # decoder layer\n",
    "        self.decoder = Transformer_decoder_net(dim_emb, nb_heads, nb_layers_decoder)\n",
    "        self.WK_att_decoder = nn.Linear(dim_emb, nb_layers_decoder* dim_emb) \n",
    "        self.WV_att_decoder = nn.Linear(dim_emb, nb_layers_decoder* dim_emb) \n",
    "        self.PE = generate_positional_encoding(dim_emb, max_len_PE)        \n",
    "        \n",
    "    def forward(self, x, B, greedy, beamsearch):\n",
    "        \n",
    "        # some parameters\n",
    "        bsz = x.shape[0]\n",
    "        nb_nodes = x.shape[1]\n",
    "        zero_to_bsz = torch.arange(bsz, device=x.device) # [0,1,...,bsz-1]\n",
    "        \n",
    "        # For beam search\n",
    "        zero_to_B = torch.arange(B, device=x.device) # [0,1,...,B-1]\n",
    "\n",
    "        # input embedding layer\n",
    "        h = self.input_emb(x) # size(h)=(bsz, nb_nodes, dim_emb)\n",
    "        \n",
    "        # concat the nodes and the input placeholder that starts the decoding\n",
    "        h = torch.cat([h, self.start_placeholder.repeat(bsz, 1, 1)], dim=1) # size(start_placeholder)=(bsz, nb_nodes+1, dim_emb)\n",
    "        \n",
    "        # encoder layer\n",
    "        h_encoder, _ = self.encoder(h) # size(h)=(bsz, nb_nodes+1, dim_emb)\n",
    "\n",
    "        # key and value for decoder    \n",
    "        K_att_decoder = self.WK_att_decoder(h_encoder) # size(K_att)=(bsz, nb_nodes+1, dim_emb*nb_layers_decoder)\n",
    "        V_att_decoder = self.WV_att_decoder(h_encoder) # size(V_att)=(bsz, nb_nodes+1, dim_emb*nb_layers_decoder)\n",
    "        \n",
    "        # starting node in tour\n",
    "        self.PE = self.PE.to(x.device)\n",
    "        \n",
    "        # For beam search\n",
    "        tours_greedy = torch.zeros(2, nb_nodes, device=x.device)\n",
    "        tours_beamsearch = torch.zeros(2, nb_nodes, device=x.device)\n",
    "        scores_greedy = torch.zeros(2, device=x.device)\n",
    "        scores_beamsearch = torch.zeros(2, device=x.device)\n",
    "        \n",
    "        # Greedy search\n",
    "        if greedy:\n",
    "            #print('Greedy decoding')\n",
    "            deterministic = True\n",
    "            # list that will contain Long tensors of shape (bsz,) that gives the idx of the cities chosen at time t\n",
    "            tours = []\n",
    "            # list that will contain Float tensors of shape (bsz,) that gives the neg log probs of the choices made at time t\n",
    "            sumLogProbOfActions = []\n",
    "            # input placeholder that starts the decoding\n",
    "            idx_start_placeholder = torch.Tensor([nb_nodes]).long().repeat(bsz).to(x.device)\n",
    "            h_start = h_encoder[zero_to_bsz, idx_start_placeholder, :] + self.PE[0].repeat(bsz,1) # size(h_start)=(bsz, dim_emb)\n",
    "            # initialize mask of visited cities\n",
    "            mask_visited_nodes = torch.zeros(bsz, nb_nodes+1, device=x.device).bool() # False\n",
    "            mask_visited_nodes[zero_to_bsz, idx_start_placeholder] = True\n",
    "            # clear key and val stored in the decoder\n",
    "            self.decoder.reset_selfatt_keys_values()\n",
    "            # construct tour recursively\n",
    "            h_t = h_start\n",
    "            for t in range(nb_nodes):\n",
    "                # compute probability over the next node in the tour\n",
    "                prob_next_node = self.decoder(h_t, K_att_decoder, V_att_decoder, mask_visited_nodes) # size(prob_next_node)=(bsz, nb_nodes+1)\n",
    "                # choose node with highest probability or sample with Bernouilli \n",
    "                if deterministic:\n",
    "                    idx = torch.argmax(prob_next_node, dim=1) # size(query)=(bsz,)\n",
    "                else:\n",
    "                    idx = Categorical(prob_next_node).sample() # size(query)=(bsz,)\n",
    "                # compute logprobs of the action items in the list sumLogProbOfActions   \n",
    "                ProbOfChoices = prob_next_node[zero_to_bsz, idx] \n",
    "                sumLogProbOfActions.append( torch.log(ProbOfChoices) )  # size(query)=(bsz,)\n",
    "                # update embedding of the current visited node\n",
    "                h_t = h_encoder[zero_to_bsz, idx, :] # size(h_start)=(bsz, dim_emb)\n",
    "                h_t = h_t + self.PE[t+1].expand(bsz, self.dim_emb)\n",
    "                # update tour\n",
    "                tours.append(idx)\n",
    "                # update masks with visited nodes\n",
    "                mask_visited_nodes = mask_visited_nodes.clone()\n",
    "                mask_visited_nodes[zero_to_bsz, idx] = True\n",
    "            # logprob_of_choices = sum_t log prob( pi_t | pi_(t-1),...,pi_0 )\n",
    "            sumLogProbOfActions = torch.stack(sumLogProbOfActions,dim=1).sum(dim=1) # size(sumLogProbOfActions)=(bsz,)\n",
    "            # convert the list of nodes into a tensor of shape (bsz,num_cities)\n",
    "            tours = torch.stack(tours,dim=1) # size(col_index)=(bsz, nb_nodes)\n",
    "            tours_greedy = tours\n",
    "            scores_greedy = sumLogProbOfActions \n",
    "        \n",
    "        # Beamsearch\n",
    "        if beamsearch:\n",
    "            #print('Beam search decoding')\n",
    "            # clear key and val stored in the decoder\n",
    "            self.decoder.reset_selfatt_keys_values() \n",
    "            K_att_decoder_tmp = K_att_decoder # size(K_att_decoder_tmp)=(bsz, nb_nodes, dim_emb*nb_layers_decoder)\n",
    "            V_att_decoder_tmp = V_att_decoder # size(V_att_decoder_tmp)=(bsz, nb_nodes, dim_emb*nb_layers_decoder)\n",
    "            for t in range(nb_nodes):\n",
    "                #if not t%10:\n",
    "                #    print('t: {}, GPU reserved mem: {:.2f}, GPU allocated mem: {:.2f}'.format(t,torch.cuda.memory_reserved(0)/1e9,torch.cuda.memory_allocated(0)/1e9))\n",
    "                if t==0: # at t=0, there are at most B_{t=0}=nb_nodes beams\n",
    "                    B_t0 = min(B, nb_nodes)\n",
    "                    # input placeholder that starts the decoding\n",
    "                    idx_start_placeholder = torch.Tensor([nb_nodes]).long().repeat(bsz).to(x.device)\n",
    "                    h_start = h_encoder[zero_to_bsz, idx_start_placeholder, :] + self.PE[0].repeat(bsz,1) # size(h_start)=(bsz, dim_emb)\n",
    "                    h_t = h_start # size(h_start)=(bsz, dim_emb)\n",
    "                    mask_visited_nodes = torch.zeros(bsz, nb_nodes+1, device=x.device).bool() # False, size(mask_visited_nodes)=(bsz, nb_nodes+1) # initialize mask of visited cities\n",
    "                    mask_visited_nodes[zero_to_bsz, idx_start_placeholder] = True\n",
    "                    # compute probability over the next node in the tour\n",
    "                    prob_next_node = self.decoder(h_t, K_att_decoder, V_att_decoder, mask_visited_nodes) # size(prob_next_node)=(bsz, nb_nodes+1) \n",
    "                    # compute score_t + sum_t score_{t-1} for all beams\n",
    "                    score_t = torch.log(prob_next_node) # size(score_t)=(bsz, nb_nodes+1) for t=0 \n",
    "                    sum_scores = score_t # size(score_t)=(bsz, nb_nodes+1)\n",
    "                    # choose nodes with top-B sumScores \n",
    "                    top_val, top_idx = torch.topk(sum_scores, B_t0, dim=1) # size(sumScores)=(bsz, B_t0)\n",
    "                    # update sum_t score_{t} for all beams\n",
    "                    sum_scores = top_val # size(sumScores)=(bsz, B_t0) \n",
    "                    zero_to_B_t0 = torch.arange(B_t0, device=x.device) # [0,1,...,B_t0-1]\n",
    "                    mask_visited_nodes = mask_visited_nodes.unsqueeze(1) # size(mask_visited_nodes)=(bsz, 1, nb_nodes+1)\n",
    "                    mask_visited_nodes = torch.repeat_interleave(mask_visited_nodes, B_t0, dim=1)\n",
    "                    for b in range(bsz):\n",
    "                        mask_visited_nodes[b, zero_to_B_t0, top_idx[b]] = True # size(mask_visited_nodes)=(bsz, B_t0, nb_nodes+1)\n",
    "                    tours = torch.zeros(bsz, B_t0, nb_nodes, device=x.device).long() # size(tours)=(bsz, B_t0, nb_nodes)\n",
    "                    tours[:,:,t] = top_idx # size(tours)=(bsz, B_t0, nb_nodes)\n",
    "                    # update embedding of the current visited node\n",
    "                    h_t = torch.zeros(bsz, B_t0, self.dim_emb, device=x.device) # size(tours)=(bsz, B_t0, dim_emb)\n",
    "                    for b in range(bsz):\n",
    "                        h_t[b, :, :] = h_encoder[b, top_idx[b], :] # size(h_t)=(bsz, B_t0, dim_emb)\n",
    "                    h_t = h_t + self.PE[t+1].expand(bsz, B_t0, self.dim_emb) # size(h_t)=(bsz, B_t0, dim_emb)\n",
    "                    self.decoder.repeat_selfatt_keys_values(B_t0)\n",
    "                    K_att_decoder = torch.repeat_interleave(K_att_decoder_tmp, B_t0, dim=0) # size(K_att_decoder)=(bsz*B_t0, nb_nodes+1, dim_emb*nb_layers_decoder)\n",
    "                    V_att_decoder = torch.repeat_interleave(V_att_decoder_tmp, B_t0, dim=0) # size(V_att_decoder)=(bsz*B_t0, nb_nodes+1, dim_emb*nb_layers_decoder)\n",
    "                    \n",
    "                elif t==1: # at t=1, there are at most B_{t=1}=nb_nodes^2 beams\n",
    "                    # compute probability over the next node in the tour\n",
    "                    h_t = h_t.view(bsz*B_t0, self.dim_emb)\n",
    "                    mask_visited_nodes = mask_visited_nodes.view(bsz*B_t0, nb_nodes+1)\n",
    "                    prob_next_node = self.decoder(h_t, K_att_decoder, V_att_decoder, mask_visited_nodes) # size(prob_next_node)=(bsz.B_t0, nb_nodes+1) \n",
    "                    prob_next_node = prob_next_node.view(bsz, B_t0, nb_nodes+1) # size(prob_next_node)=(bsz, B_t0, nb_nodes+1) \n",
    "                    mask_visited_nodes = mask_visited_nodes.view(bsz, B_t0, nb_nodes+1)\n",
    "                    h_t = h_t.view(bsz, B_t0, self.dim_emb) \n",
    "                    # compute score_t + sum_t score_{t-1} for all beams\n",
    "                    score_t = torch.log(prob_next_node) # size(score_t)=(bsz, B, nb_nodes+1) \n",
    "                    sum_scores = score_t + sum_scores.unsqueeze(2) # size(score_t)=(bsz, B, nb_nodes+1)\n",
    "                    sum_scores_flatten = sum_scores.view(bsz, -1) # size(sumScores_next_node)=(bsz, B.(nb_nodes+1))\n",
    "                    # choose nodes with top-B sumScores \n",
    "                    top_val, top_idx = torch.topk(sum_scores_flatten, B, dim=1)\n",
    "                    idx_top_beams = top_idx // (nb_nodes+1) # size(idx_beam_topB)=(bsz, B)\n",
    "                    idx_in_beams = top_idx - idx_top_beams* (nb_nodes+1) # size(idx_in_beams)=(bsz, B)\n",
    "                    # update sum_t score_{t} for all beams\n",
    "                    sum_scores = top_val\n",
    "                    # update beam masks with visited nodes\n",
    "                    mask_visited_nodes_tmp = mask_visited_nodes.clone() # size(mask_visited_nodes_tmp)=(bsz, B_t0, nb_nodes+1)\n",
    "                    mask_visited_nodes = torch.zeros(bsz, B, nb_nodes+1, device=x.device).bool() # size(mask_visited_nodes)=(bsz, B, nb_nodes+1)\n",
    "                    for b in range(bsz):\n",
    "                        mask_visited_nodes[b, zero_to_B, :] = mask_visited_nodes_tmp[b, idx_top_beams[b], :] # size(mask_visited_nodes)=(bsz, B, nb_nodes+1)\n",
    "                    for b in range(bsz):\n",
    "                        mask_visited_nodes[b, zero_to_B, idx_in_beams[b]] = True # size(mask_visited_nodes)=(bsz, B, nb_nodes+1)\n",
    "                    # update beam tours with visited nodes\n",
    "                    tours_tmp = tours.clone()\n",
    "                    tours = torch.zeros(bsz, B, nb_nodes, device=x.device).long() # size(tours)=(bsz, B, nb_nodes)\n",
    "                    for b in range(bsz):\n",
    "                        tours[b, zero_to_B, :] = tours_tmp[b, idx_top_beams[b], :]\n",
    "                    tours[:,:,t] = idx_in_beams # size(tours)=(bsz, B, nb_nodes)\n",
    "                    # update embedding of the current visited node\n",
    "                    h_t = torch.zeros(bsz, B, self.dim_emb, device=x.device) # size(tours)=(bsz, B_t0, dim_emb)\n",
    "                    for b in range(bsz):\n",
    "                        h_t[b, :, :] = h_encoder[b, idx_in_beams[b], :] # size(h_t)=(bsz, B, dim_emb)\n",
    "                    h_t = h_t + self.PE[t+1].expand(bsz, B, self.dim_emb) # size(h_t)=(bsz, B, dim_emb)\n",
    "                    # update self-attention embeddings of partial tours\n",
    "                    self.decoder.reorder_selfatt_keys_values(t, idx_top_beams) # size(K_att_decoder)=(bsz*B_t0, nb_nodes+1, dim_emb*nb_layers_decoder)\n",
    "                    K_att_decoder = torch.repeat_interleave(K_att_decoder_tmp, B, dim=0) # size(K_att_decoder)=(bsz*B, nb_nodes+1, dim_emb*nb_layers_decoder)\n",
    "                    V_att_decoder = torch.repeat_interleave(V_att_decoder_tmp, B, dim=0) # size(V_att_decoder)=(bsz*B, nb_nodes+1, dim_emb*nb_layers_decoder)\n",
    "\n",
    "                else: # at t>=2, we arbitrary decide to have at most B_{t>=2}=nb_nodes^2 beams\n",
    "                    # compute probability over the next node in the tour\n",
    "                    h_t = h_t.view(bsz*B, self.dim_emb)\n",
    "                    mask_visited_nodes = mask_visited_nodes.view(bsz*B, nb_nodes+1)\n",
    "                    prob_next_node = self.decoder(h_t, K_att_decoder, V_att_decoder, mask_visited_nodes) # size(prob_next_node)=(bsz.B, nb_nodes+1) \n",
    "                    prob_next_node = prob_next_node.view(bsz, B, nb_nodes+1) # size(prob_next_node)=(bsz, B, nb_nodes+1) \n",
    "                    mask_visited_nodes = mask_visited_nodes.view(bsz, B, nb_nodes+1)\n",
    "                    h_t = h_t.view(bsz, B, self.dim_emb) \n",
    "                    # compute score_t + sum_t score_{t-1} for all beams\n",
    "                    score_t = torch.log(prob_next_node) # size(score_t)=(bsz, B, nb_nodes+1)\n",
    "                    sum_scores = score_t + sum_scores.unsqueeze(2) # size(score_t)=(bsz, B, nb_nodes+1)\n",
    "                    sum_scores_flatten = sum_scores.view(bsz, -1) # size(sumScores_next_node)=(bsz, B.(nb_nodes+1))\n",
    "                    # choose nodes with top-B sumScores \n",
    "                    top_val, top_idx = torch.topk(sum_scores_flatten, B, dim=1)\n",
    "                    idx_top_beams = top_idx // (nb_nodes+1) # size(idx_beam_topB)=(bsz, B)\n",
    "                    idx_in_beams = top_idx - idx_top_beams* (nb_nodes+1) # size(idx_in_beams)=(bsz, B)\n",
    "                    # update sum_t score_{t} for all beams\n",
    "                    sum_scores = top_val\n",
    "                    # update beam masks with visited nodes\n",
    "                    mask_visited_nodes_tmp = mask_visited_nodes.clone()\n",
    "                    for b in range(bsz):\n",
    "                        mask_visited_nodes[b, zero_to_B, :] = mask_visited_nodes_tmp[b, idx_top_beams[b], :]\n",
    "                    for b in range(bsz):\n",
    "                        mask_visited_nodes[b, zero_to_B, idx_in_beams[b]] = True\n",
    "                    # update beam tours with visited nodes\n",
    "                    tours_tmp = tours.clone()\n",
    "                    for b in range(bsz):\n",
    "                        tours[b, zero_to_B, :] = tours_tmp[b, idx_top_beams[b], :]\n",
    "                    tours[:,:,t] = idx_in_beams # size(tours)=(bsz, B, nb_nodes)\n",
    "                    # update embedding of the current visited node\n",
    "                    for b in range(bsz):\n",
    "                        h_t[b, :, :] = h_encoder[b, idx_in_beams[b], :] # size(h_t)=(bsz, B, dim_emb)\n",
    "                    h_t = h_t + self.PE[t+1].expand(bsz, B, self.dim_emb) # size(h_t)=(bsz, B, dim_emb)\n",
    "                    # update self-attention embeddings of partial tours\n",
    "                    self.decoder.reorder_selfatt_keys_values(t, idx_top_beams)\n",
    "            # sum_t log prob( pi_t | pi_0,...pi_(t-1) )\n",
    "            sum_scores = sum_scores[:,0] # size(sumScores)=(bsz)\n",
    "            tours_beamsearch = tours\n",
    "            scores_beamsearch = sum_scores\n",
    "\n",
    "        return tours_greedy, tours_beamsearch, scores_greedy, scores_beamsearch\n",
    "    \n",
    "    \n",
    "    \n",
    "\n",
    "###################\n",
    "# Instantiate a training network and a baseline network\n",
    "###################\n",
    "try: \n",
    "    del model_baseline # remove existing model\n",
    "except:\n",
    "    pass\n",
    "\n",
    "model_baseline = TSP_net(args.dim_input_nodes, args.dim_emb, args.dim_ff, \n",
    "              args.nb_layers_encoder, args.nb_layers_decoder, args.nb_heads, args.max_len_PE,\n",
    "              batchnorm=args.batchnorm)\n",
    "\n",
    "# # uncomment these lines if trained with multiple GPUs\n",
    "# print(torch.cuda.device_count())\n",
    "# if torch.cuda.device_count()>1:\n",
    "#     model_baseline = nn.DataParallel(model_baseline)\n",
    "# # uncomment these lines if trained with multiple GPUs\n",
    "\n",
    "model_baseline = model_baseline.to(device)\n",
    "model_baseline.eval()\n",
    "\n",
    "print(args); print('')\n",
    "\n",
    "\n",
    "\n",
    "###################\n",
    "# Load checkpoint\n",
    "###################\n",
    "checkpoint_file = \"checkpoint/checkpoint_21-03-01--17-25-00-n50-gpu0.pkl\" \n",
    "checkpoint = torch.load(checkpoint_file, map_location=device)\n",
    "epoch_ckpt = checkpoint['epoch'] + 1\n",
    "tot_time_ckpt = checkpoint['tot_time']\n",
    "plot_performance_train = checkpoint['plot_performance_train']\n",
    "plot_performance_baseline = checkpoint['plot_performance_baseline']\n",
    "model_baseline.load_state_dict(checkpoint['model_baseline'])\n",
    "print('Load checkpoint file={:s}\\n  Checkpoint at epoch= {:d} and time={:.3f}min\\n'.format(checkpoint_file,epoch_ckpt-1,tot_time_ckpt/60))\n",
    "del checkpoint\n",
    "\n",
    "mystring_min = 'Epoch: {:d}, tot_time_ckpt: {:.3f}day, L_train: {:.3f}, L_base: {:.3f}\\n'.format(\n",
    "    epoch_ckpt, tot_time_ckpt/3660/24, plot_performance_train[-1][1], plot_performance_baseline[-1][1]) \n",
    "print(mystring_min) \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "###################   \n",
    "# Hyper-parameter for beam search\n",
    "###################\n",
    "\n",
    "# experiment n50\n",
    "args.nb_nodes = 50\n",
    "# B = 1; args.bsz = 10000; greedy = True; beamsearch = False # greedy\n",
    "# B = 100; args.bsz = 250; greedy = False; beamsearch = True \n",
    "# B = 1000; args.bsz = 25; greedy = False; beamsearch = True \n",
    "B = 2500; args.bsz = 10; greedy = False; beamsearch = True \n",
    "args.nb_batch_eval = (10000+args.bsz-1)// args.bsz\n",
    "print('nb_nodes: {}, bsz: {}, B: {}, nb_batch_eval: {}, tot TSPs: {}\\n'.format(args.nb_nodes, args.bsz, B, args.nb_batch_eval, args.nb_batch_eval* args.bsz))\n",
    "\n",
    "# # experiment n100\n",
    "# args.nb_nodes = 100\n",
    "# B = 1; args.bsz = 5000; greedy = True; beamsearch = False # greedy\n",
    "# B = 1000; args.bsz = 10; greedy = False; beamsearch = True \n",
    "# args.nb_batch_eval = (10000+args.bsz-1)// args.bsz\n",
    "# print('nb_nodes: {}, bsz: {}, B: {}, nb_batch_eval: {}, tot TSPs: {}\\n'.format(args.nb_nodes, args.bsz, B, args.nb_batch_eval, args.nb_batch_eval* args.bsz))\n",
    "\n",
    "\n",
    "\n",
    "###################   \n",
    "# Test set\n",
    "###################\n",
    "if args.nb_nodes == 50:\n",
    "    x_10k = torch.load('data/10k_TSP50.pt').to(device)\n",
    "    x_10k_len = torch.load('data/10k_TSP50_len.pt').to(device)\n",
    "    L_concorde = x_10k_len.mean().item()\n",
    "if args.nb_nodes == 100:\n",
    "    x_10k = torch.load('data/10k_TSP100.pt').to(device)\n",
    "    x_10k_len = torch.load('data/10k_TSP100_len.pt').to(device)\n",
    "    L_concorde = x_10k_len.mean().item()\n",
    "nb_TSPs = args.nb_batch_eval* args.bsz\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "    \n",
    "    \n",
    "###################   \n",
    "# Run beam search\n",
    "###################\n",
    "start = time.time()\n",
    "mean_tour_length_greedy = 0\n",
    "mean_tour_length_beamsearch = 0\n",
    "mean_scores_greedy = 0\n",
    "mean_scores_beamsearch = 0\n",
    "gap_greedy = 0\n",
    "gap_beamsearch = 0\n",
    "for step in range(0,args.nb_batch_eval):\n",
    "    print('batch index: {}, tot_time: {:.3f}min'.format(step, (time.time()-start)/60))\n",
    "    # extract a batch of test tsp instances \n",
    "    x = x_10k[step*args.bsz:(step+1)*args.bsz,:,:]\n",
    "    x_len_concorde = x_10k_len[step*args.bsz:(step+1)*args.bsz]\n",
    "    # compute tour for model and baseline\n",
    "    with torch.no_grad():\n",
    "        tours_greedy, tours_beamsearch, scores_greedy, scores_beamsearch = model_baseline(x, B, greedy, beamsearch)\n",
    "        # greedy\n",
    "        if greedy:\n",
    "            L_greedy = compute_tour_length(x, tours_greedy)\n",
    "            mean_tour_length_greedy += L_greedy.mean().item()  \n",
    "            mean_scores_greedy += scores_greedy.mean().item()  \n",
    "            x_len_greedy = L_greedy\n",
    "            gap_greedy += (x_len_greedy/ x_len_concorde - 1.0).sum()   \n",
    "        # beamsearch\n",
    "        if beamsearch:\n",
    "            tours_beamsearch = tours_beamsearch.view(args.bsz*B, args.nb_nodes)\n",
    "            x = x.repeat_interleave(B,dim=0)\n",
    "            L_beamsearch = compute_tour_length(x, tours_beamsearch)\n",
    "            tours_beamsearch = tours_beamsearch.view(args.bsz, B, args.nb_nodes)\n",
    "            L_beamsearch = L_beamsearch.view(args.bsz, B)\n",
    "            L_beamsearch_tmp = L_beamsearch\n",
    "            L_beamsearch, idx_min = L_beamsearch.min(dim=1)\n",
    "            mean_tour_length_beamsearch += L_beamsearch.mean().item()\n",
    "            mean_scores_beamsearch += scores_beamsearch.mean().item()  \n",
    "            x_len_beamsearch = L_beamsearch\n",
    "            gap_beamsearch += (x_len_beamsearch/ x_len_concorde - 1.0).sum()\n",
    "    torch.cuda.empty_cache() # free GPU reserved memory \n",
    "if greedy:\n",
    "    mean_tour_length_greedy =  mean_tour_length_greedy/ args.nb_batch_eval\n",
    "    mean_scores_greedy =  mean_scores_greedy/ args.nb_batch_eval\n",
    "    gap_greedy = (gap_greedy/ nb_TSPs).item()\n",
    "if beamsearch:\n",
    "    mean_tour_length_beamsearch =  mean_tour_length_beamsearch/ args.nb_batch_eval\n",
    "    mean_scores_beamsearch =  mean_scores_beamsearch/ args.nb_batch_eval\n",
    "    gap_beamsearch /= nb_TSPs\n",
    "tot_time = time.time()-start\n",
    "    \n",
    "\n",
    "    \n",
    "###################   \n",
    "# Write result file\n",
    "###################\n",
    "nb_TSPs = args.nb_batch_eval* args.bsz\n",
    "file_name = \"beamsearch-nb_nodes{}\".format(args.nb_nodes) + \"-nb_TSPs{}\".format(nb_TSPs) + \"-B{}\".format(B) + \".txt\"\n",
    "file = open(file_name,\"w\",1) \n",
    "mystring = '\\nnb_nodes: {:d}, nb_TSPs: {:d}, B: {:d}, L_greedy: {:.6f}, L_concorde: {:.5f}, L_beamsearch: {:.5f}, \\\n",
    "gap_greedy(%): {:.5f}, gap_beamsearch(%): {:.5f}, scores_greedy: {:.5f}, scores_beamsearch: {:.5f}, tot_time: {:.4f}min, \\\n",
    "tot_time: {:.3f}hr, mean_time: {:.3f}sec'.format(args.nb_nodes, nb_TSPs, B, mean_tour_length_greedy, L_concorde, \\\n",
    "                                 mean_tour_length_beamsearch, 100*gap_greedy, 100*gap_beamsearch, mean_scores_greedy, \\\n",
    "                                 mean_scores_beamsearch, tot_time/60, tot_time/3600, tot_time/nb_TSPs)\n",
    "print(mystring)\n",
    "file.write(mystring)\n",
    "file.close()\n"
   ]
  },
  {
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}