Second commit... full folder...

src/vrp/Baseline.py

@@ -0,0 +1,37 @@
+import torch
+
+from Model import AttentionModel
+
+
+class Baseline(object):
+    def __init__(self, model: AttentionModel, graph_size, nb_val_samples):
+        self.model = model
+        self.graph_size = graph_size
+        self.nb_val_samples = nb_val_samples
+        self.tour_lengths = []
+
+    def eval(self):
+        with torch.no_grad():
+            return self.model.eval()
+
+    def epoch_callback(self, model, epoch):
+        pass
+
+    def state_dict(self):
+        return self.model.state_dict()
+
+    def load_state_dict(self, state_dict):
+        return self.model.load_state_dict(state_dict)
+
+    def evaluate(self, inputs, use_solver=False):
+        if use_solver:
+            raise NotImplementedError("Solver doesnt have a stable implementation yet.")
+        _, _, rollout_sol = self.model(inputs)
+        self.tour_lengths = self.model.compute(inputs, rollout_sol)
+        return self.tour_lengths
+
+    def set_decode_mode(self, mode):
+        self.model.set_decode_mode(mode)
+
+    def reset(self):
+        self.tour_lengths = []

src/vrp/Decoder.py

@@ -0,0 +1,128 @@
+from math import sqrt
+
+import torch
+import torch.nn as nn
+from accelerate import Accelerator
+from torch.distributions import Categorical
+
+from Encoder import MultiHeadAttention
+
+
+class Decoder(nn.Module):
+    """
+        This class contains the decoder that will be used to compute the probability distribution from which we will sample
+        which city to visit next.
+    """
+
+    def __init__(self, n_head, embedding_dim, decode_mode="sample", C=10):
+        super(Decoder, self).__init__()
+        self.scale = sqrt(embedding_dim)
+        self.decode_mode = decode_mode
+        self.C = C
+
+        self.vl = nn.Parameter(
+            torch.FloatTensor(size=[1, 1, embedding_dim]).uniform_(-1. / embedding_dim, 1. / embedding_dim),
+            requires_grad=True)
+        self.vf = nn.Parameter(
+            torch.FloatTensor(size=[1, 1, embedding_dim]).uniform_(-1. / embedding_dim, 1. / embedding_dim),
+            requires_grad=True)
+
+        self.glimpse = MultiHeadAttention(n_head, embedding_dim, 3 * embedding_dim, embedding_dim, embedding_dim)
+        self.project_k = nn.Linear(embedding_dim, embedding_dim, bias=False)
+        self.cross_entropy = nn.CrossEntropyLoss(reduction='none')
+
+        self.accelerate = Accelerator()
+        self.glimpse, self.project_k, self.cross_entropy = self.accelerate.prepare(self.glimpse, self.project_k, self.cross_entropy)
+
+    def forward(self, inputs):
+        """
+
+        :param inputs: (encoded_inputs, demands, capacities) ([batch_size, seq_len, embedding_dim],[batch_size, seq_len],[batch_size])
+        :return: log_prob, solutions
+        """
+        encoded_inputs, demands, capacities = inputs
+
+        batch_size, seq_len, embedding_dim = encoded_inputs.size()  # sel_len = nb_clients + nb_depot (1)
+
+        h_hat = encoded_inputs.mean(-2, keepdim=True)
+
+        city_index = None
+
+        # case of vrp : mask[:, 0] is the depot
+        mask = torch.zeros([batch_size, seq_len], device=self.accelerate.device).bool()
+
+        solution = torch.zeros([batch_size, 1], dtype=torch.long, device=self.accelerate.device)  # first node is depot
+        mask[:, 0] = True  # vehicle is in the depot location
+
+        log_probabilities = torch.zeros(batch_size, dtype=torch.float32, device=self.accelerate.device)
+
+        last = self.vl.repeat(batch_size, 1, 1)  # batch_size, 1, embedding_dim
+
+        first = self.vf.repeat(batch_size, 1, 1)  # batch_size, 1, embedding_dim
+
+        raw_logits = torch.tensor([], device=self.accelerate.device)
+        t = 0  # time steps
+
+        # for t in range(seq_len):
+        while torch.sum(mask) < batch_size * seq_len:
+            t += 1
+            h_c = torch.cat((h_hat, last, first), dim=-1)  # [batch_size, 1, 3 * embedding_size]
+
+            context = self.glimpse(h_c, encoded_inputs, encoded_inputs,
+                                   mask.unsqueeze(1).unsqueeze(1))
+
+            k = self.project_k(encoded_inputs)
+
+            u = torch.tanh(torch.matmul(context, k.clone().transpose(-2, -1)) / self.scale) * self.C
+
+            raw_logits = torch.cat((raw_logits, u), dim=1)
+
+            u = u.masked_fill(mask.unsqueeze(1), float('-inf'))
+
+            probas = nn.functional.softmax(u.squeeze(1), dim=-1)
+
+            one_hot = torch.zeros([seq_len], device=self.accelerate.device)
+            one_hot[0] = 1
+
+            if self.decode_mode == "greedy":
+                proba, city_index = self.greedy_decoding(probas)
+            elif self.decode_mode == "sample":
+                proba, city_index = self.sample_decoding(probas)
+
+            log_probabilities += self.cross_entropy(u.squeeze(1), city_index.view(-1))
+
+            solution = torch.cat((solution, city_index), dim=1)
+
+            # next node for the query
+            last = encoded_inputs[[i for i in range(batch_size)], city_index.view(-1), :].unsqueeze(1)
+
+            # update mask
+            mask = mask.scatter(1, city_index, True)
+
+            if t == 1:
+                first = last
+        return raw_logits, log_probabilities, solution
+
+    @staticmethod
+    def greedy_decoding(probas):
+        """
+        :param probas: [batch_size, seq_len]
+        :return: probas : [batch_size],  city_index: [batch_size,1]
+        """
+        probas, city_index = torch.max(probas, dim=1)
+
+        return probas, city_index.view(-1, 1)
+
+    @staticmethod
+    def sample_decoding(probas):
+        """
+
+        :param probas: [ batch_size, seq_len]
+        :return: probas : [batch_size],  city_index: [batch_size,1]
+        """
+        batch_size = probas.size(0)
+        m = Categorical(probas)
+        city_index = m.sample()
+        probas = probas[[i for i in range(batch_size)], city_index]
+
+        return probas, city_index.view(-1, 1)

src/vrp/Encoder.py

@@ -0,0 +1,116 @@
+from math import sqrt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as functional
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, n_heads, embedding_dim, q_dim, k_dim, v_dim):
+        super(MultiHeadAttention, self).__init__()
+
+        self.n_head = n_heads
+        self.hidden_dim = embedding_dim // self.n_head
+
+        self.queries_transformation = nn.Linear(q_dim, embedding_dim, bias=False)
+        self.keys_transformation = nn.Linear(k_dim, embedding_dim, bias=False)
+        self.values_transformation = nn.Linear(v_dim, embedding_dim, bias=False)
+
+        self.out = nn.Linear(embedding_dim, embedding_dim, bias=False)
+
+    def forward(self, query, key, value, mask=None):
+        """
+        :param query: [batch_size]
+        :param key: [seq_len]
+        :param value: [embedding_dim]
+        :param mask : [batch_size, seq_len]
+        :return: out : [batch_size, seq_len, embedding_dim]
+        """
+        batch_size = query.size(0)
+
+        q = self.queries_transformation(query).reshape(batch_size, -1, self.n_head, self.hidden_dim).permute(0, 2, 1, 3)
+        k = self.keys_transformation(key).reshape(batch_size, -1, self.n_head, self.hidden_dim).permute(0, 2, 1, 3)
+        v = self.values_transformation(value).reshape(batch_size, -1, self.n_head, self.hidden_dim).permute(0, 2, 1, 3)
+
+        scores = torch.matmul(q, k.transpose(3, 2)) / sqrt(self.hidden_dim)
+
+        if mask is not None:
+            scores = scores.masked_fill(mask, float('-inf'))
+
+        attention = torch.nn.functional.softmax(scores, dim=-1)
+
+        output = torch.matmul(attention, v)
+
+        concat_output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.n_head * self.hidden_dim)
+
+        out = self.out(concat_output)
+
+        return out
+
+
+class FeedForwardLayer(nn.Module):
+    def __init__(self, embedding_dim, dim_feedforward):
+        super(FeedForwardLayer, self).__init__()
+        self.embedding_dim = embedding_dim
+        self.dim_feedforward = dim_feedforward
+
+        self.lin1 = nn.Linear(self.embedding_dim, self.dim_feedforward)
+        self.lin2 = nn.Linear(self.dim_feedforward, self.embedding_dim)
+
+    def forward(self, inputs):
+        """
+
+        :param inputs: [batch_size, seq_len, embedding_dim]
+        :return: out : [batch_size, seq_len, embedding_dim]
+        """
+        return self.lin2(functional.relu(self.lin1(inputs)))
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, n_head, embedding_dim, dim_feedforward, dropout=0.1):
+        super(EncoderLayer, self).__init__()
+
+        self.multi_head_attention = MultiHeadAttention(n_head,
+                                                       embedding_dim,
+                                                       embedding_dim,
+                                                       embedding_dim,
+                                                       embedding_dim)
+        self.feed_forward_layer = FeedForwardLayer(embedding_dim, dim_feedforward)
+
+        self.dropout1 = nn.Dropout(dropout)  # MHA dropout
+        self.dropout2 = nn.Dropout(dropout)  # FFL dropout
+
+        self.bn1 = nn.BatchNorm1d(embedding_dim, affine=True)
+        self.bn2 = nn.BatchNorm1d(embedding_dim, affine=True)
+
+    def forward(self, x):
+        """
+
+        :param x: [batch_size, seq_len, embedding_dim]
+        :return: out : [batch_size, seq_len, embedding_dim]
+        """
+
+        x = x + self.dropout1(self.multi_head_attention(x, x, x))
+        x = self.bn1(x.view(-1, x.size(-1))).view(*x.size())
+
+        x = x + self.dropout2(self.feed_forward_layer(x))
+        x = self.bn2(x.view(-1, x.size(-1))).view(*x.size())
+
+        return x
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(self, n_layers, n_head, embedding_dim, dim_feedforward, dropout=0.1):
+        super(TransformerEncoder, self).__init__()
+        self.layers = [EncoderLayer(n_head, embedding_dim, dim_feedforward, dropout) for _ in
+                       range(n_layers)]
+        self.transformer_encoder = nn.Sequential(*self.layers)
+
+    def forward(self, inputs):
+        """
+
+        :param inputs: [batch_size, seq_len, embedding_dim]
+        :return: [batch_size, seq_len, embedding_dim]
+        """
+
+        return self.transformer_encoder(inputs)

src/vrp/Model.py

@@ -0,0 +1,137 @@
+import time
+
+import torch
+import torch.nn as nn
+from accelerate import Accelerator
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from Decoder import Decoder
+from Encoder import TransformerEncoder
+
+
+class AttentionModel(nn.Module):
+    def __init__(self, embedding_dim, n_layers, n_head, dim_feedforward, C, dropout=0.1):
+        super(AttentionModel, self).__init__()
+        self.embedding_dim = embedding_dim
+        self.n_layers = n_layers
+        self.n_head = n_head
+        self.dim_feedforward = dim_feedforward
+        self.decode_mode = "greedy"
+        self.dropout = dropout
+        self.C = C
+        self.input_dim = 2
+
+        self.demand_embedding = nn.Linear(1, self.embedding_dim)
+
+        self.city_embedding = nn.Linear(self.input_dim, self.embedding_dim)
+
+        self.encoder = TransformerEncoder(self.n_layers, self.n_head, self.embedding_dim,
+                                          self.dim_feedforward, self.dropout)
+        self.decoder = Decoder(self.n_head, self.embedding_dim, self.decode_mode, self.C)
+
+        self.accelerator = Accelerator()
+        self.encoder, self.decoder = self.accelerator.prepare(self.encoder, self.decoder)
+
+    def forward(self, inputs):
+        """
+
+        :param inputs : (locations, demands,capacities)
+               (locations : [batch_size, seq_len, input_dim],
+                demands : [batch_size, seq_len, 1],
+                capacities : [batch_size])
+
+        :return: raw_logits : [batch_size, seq_len, seq_len],
+                 log_prob : [batch_size],
+                 solutions : [batch_size, seq_len]
+        """
+
+        inputs, demands, capacities = inputs
+        dem = demands.unsqueeze(-1)
+
+        data = self.encoder(self.city_embedding(inputs) + self.demand_embedding(dem))
+
+        raw_logits, log_prob, solution = self.decoder((data, demands, capacities))
+
+        return raw_logits, log_prob, solution
+
+    def set_decode_mode(self, mode):
+        self.decode_mode = mode
+        self.decoder.decode_mode = mode
+
+    def test(self, data: DataLoader, decode_mode="greedy"):
+        tour_lengths = torch.tensor([])
+        self.eval()
+        self.set_decode_mode(decode_mode)
+        cpu = time.time()
+
+        for batch_id, batch in enumerate(tqdm(data)):
+            locations, demands, capacities = batch
+            inputs = (locations, demands, capacities.float())
+            _, _, solution = self(inputs)
+
+            btl = self.compute(inputs, solution)
+            tour_lengths = torch.cat((tour_lengths, btl), dim=0)
+
+        cpu = time.time() - cpu
+        return {
+            "tour_lengths": tour_lengths,
+            "avg_tl": tour_lengths.mean().item(),
+            "cpu": cpu
+        }
+
+    def compute(self, instance_data, route):
+
+        batch_size = route.size(0)
+        length = torch.FloatTensor(torch.zeros(batch_size, 1, device=self.accelerator.device))
+        locations, demands, capacities = instance_data
+        for batch in range(batch_size):
+            # Ref : Thibaut Vidal, Split algorithm in O(n) for the capacitated vehicle routing problem
+            # https://arxiv.org/pdf/1508.02759.pdf
+            nb_nodes = len(route[batch]) - 1
+            distance_to_depot = torch.tensor(
+                [torch.norm(locations[batch, int(route[batch, route[batch, i]])] - locations[batch, 0]) for i in
+                 range(nb_nodes + 1)]
+            , device=self.accelerator.device)
+            distance_to_next = torch.tensor([
+                torch.sqrt(torch.sum((locations[batch, route[batch, i]] - locations[batch, route[batch, i + 1]]) ** 2))
+
+                if i < nb_nodes else -1
+                for i in range(nb_nodes + 1)], device=self.accelerator.device)
+            potential = torch.tensor([0.0] + [float('inf')] * nb_nodes, device=self.accelerator.device)
+            pred = torch.tensor([-1] * (nb_nodes + 1), device=self.accelerator.device)
+            sum_distance = torch.zeros(nb_nodes + 1, device=self.accelerator.device)
+            sum_load = torch.zeros(nb_nodes + 1, device=self.accelerator.device)
+
+            for i in range(1, nb_nodes + 1):
+                sum_load[i] = sum_load[i - 1] + demands[batch, i]
+                sum_distance[i] = sum_distance[i - 1] + distance_to_next[i - 1]
+
+            queue = [0]
+
+            for i in range(1, nb_nodes + 1):
+                potential[i] = potential[queue[0]] + sum_distance[i] - sum_distance[queue[0] + 1] + distance_to_depot[
+                    queue[0] + 1] + distance_to_depot[i]
+                pred[i] = queue[0]
+
+                if i < nb_nodes:
+                    if (sum_load[queue[-1]] != sum_load[i]) or \
+                            (potential[i] + distance_to_depot[i + 1] <= potential[queue[-1]]
+                             + distance_to_depot[queue[-1] + 1]
+                             + sum_distance[i + 1]
+                             - sum_distance[queue[-1] + 1]):
+
+                        while len(queue) > 0 and \
+                                (potential[i] + distance_to_depot[i + 1] <
+                                 potential[queue[-1]]
+                                 + distance_to_depot[queue[-1] + 1]
+                                 + sum_distance[i + 1]
+                                 - sum_distance[queue[-1] + 1]):
+                            queue.pop()
+                        queue.append(i)
+                    while sum_load[i + 1] - sum_load[queue[0]] > capacities[batch]:
+                        queue.pop(0)
+
+            length[batch] = potential[nb_nodes]
+
+        return length.view(-1)

src/vrp/Trainer.py

@@ -0,0 +1,192 @@
+import csv
+import time
+
+import numpy as np
+import torch
+from scipy.stats import ttest_rel
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+import plot
+from Baseline import Baseline
+from Model import AttentionModel
+from data.VRPDataset import VRPDataset
+from accelerate import Accelerator
+
+
+class Trainer:
+    def __init__(self, graph_size, n_epochs, batch_size, nb_train_samples,
+                 nb_val_samples, n_layers, n_heads, embedding_dim,
+                 dim_feedforward, C, dropout, learning_rate
+                 ):
+        self.n_epochs = n_epochs
+        self.batch_size = batch_size
+        self.graph_size = graph_size
+        self.nb_train_samples = nb_train_samples
+        self.nb_val_samples = nb_val_samples
+        self.accelerator = Accelerator()
+
+        self.model = AttentionModel(embedding_dim, n_layers, n_heads,
+                                    dim_feedforward,
+                                    C, dropout)  # embedding, encoder, decoder
+
+        # This ia a rollout baseline
+
+        baseline_model = AttentionModel(embedding_dim, n_layers, n_heads,
+                                        dim_feedforward,
+                                        C, dropout)
+        self.baseline = Baseline(baseline_model, graph_size, nb_val_samples)
+
+        self.baseline.load_state_dict(self.model.state_dict())
+
+        self.optimizer = Adam(self.model.parameters(), lr=learning_rate)
+
+        self.baseline, self.optimizer = self.accelerator.prepare(self.baseline, self.optimizer)
+
+        log_file_name = "{}-{}-logs.csv".format("vrp", graph_size)
+
+        f = open(log_file_name, 'w', newline='')
+        self.log_file = csv.writer(f, delimiter=",")
+
+        header = ["epoch", "losses_per_batch", "avg_tl_batch_train", "avg_tl_epoch_train", "avg_tl_epoch_val"]
+        self.log_file.writerow(header)
+
+    def train(self):
+        validation_dataset = VRPDataset(size=self.graph_size, num_samples=self.nb_val_samples)
+        print("Validation dataset created with {} samples".format(len(validation_dataset)))
+        validation_dataloader = DataLoader(validation_dataset, batch_size=self.batch_size, shuffle=True, num_workers=0,
+                                           pin_memory=True)
+        losses = []
+        avg_tour_length_batch = []
+        avg_tour_length_epoch = []
+        avg_tl_epoch_val = []
+        for epoch in range(self.n_epochs):
+            cpu = time.time()
+
+            all_tour_lengths = torch.tensor([], dtype=torch.float32, device=self.accelerator.device)
+
+            # Put model in train mode!
+            self.model.set_decode_mode("sample")
+            self.model, self.optimizer, validation_dataloader = self.accelerator.prepare(self.model, self.optimizer,
+                                                                                         validation_dataloader)
+            self.baseline.model = self.accelerator.prepare(self.baseline.model)
+            self.model.train()
+
+            # Generate new training data for each epoch
+            train_dataset = VRPDataset(size=self.graph_size, num_samples=self.nb_train_samples)
+
+            train_dataloader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=0,
+                                          pin_memory=True)
+            nb_batches = len(validation_dataloader)
+
+            for batch_id, batch in enumerate(tqdm(train_dataloader)):
+                locations, demands, capacities = batch
+
+                inputs = (locations.to(self.accelerator.device), demands.to(self.accelerator.device),
+                          capacities.float().to(self.accelerator.device))
+
+                _, log_prob, solution = self.model(inputs)
+
+                with torch.no_grad():
+                    tour_lengths = self.model.compute(inputs, solution)
+                    baseline_tour_lengths = self.baseline.evaluate(inputs, False)
+
+                    advantage = tour_lengths - baseline_tour_lengths[0:len(tour_lengths)]
+
+                loss = advantage * (-log_prob)
+                loss = loss.mean()
+
+                self.optimizer.zero_grad(set_to_none=True)
+                # loss.backward()
+                self.accelerator.backward(loss)
+                self.optimizer.step()
+
+                # save data for plot
+                losses.append(loss.item())
+                avg_tour_length_batch.append(tour_lengths.mean().item())
+                all_tour_lengths = torch.cat((all_tour_lengths, tour_lengths), dim=0)
+
+            avg_tour_length_epoch.append(all_tour_lengths.mean().item())
+
+            print(
+                "\nEpoch: {}\t\nAverage tour length model : {}\nAverage tour length baseline : {}\n".format(
+                    epoch + 1, tour_lengths.mean(), baseline_tour_lengths.mean()
+                ))
+
+            print("Validation and rollout update check\n")
+            # t-test :
+            self.model.set_decode_mode("greedy")
+            self.baseline.set_decode_mode("greedy")
+            self.model.eval()
+            self.baseline.eval()
+            with torch.no_grad():
+                rollout_tl = torch.tensor([], dtype=torch.float32)
+                policy_tl = torch.tensor([], dtype=torch.float32)
+                for batch_id, batch in enumerate(tqdm(validation_dataloader)):
+                    locations, demands, capacities = batch
+
+                    inputs = (locations, demands, capacities.float())
+
+                    _, _, solution = self.model(inputs)
+
+                    tour_lengths = self.model.compute(inputs, solution)
+                    baseline_tour_lengths = self.baseline.evaluate(inputs, False)
+
+                    rollout_tl = torch.cat((rollout_tl, baseline_tour_lengths.view(-1).cpu()), dim=0)
+                    policy_tl = torch.cat((policy_tl, tour_lengths.view(-1).cpu()), dim=0)
+
+                rollout_tl = rollout_tl.cpu().numpy()
+                policy_tl = policy_tl.cpu().numpy()
+
+                avg_ptl = np.mean(policy_tl)
+                avg_rtl = np.mean(rollout_tl)
+
+                avg_tl_epoch_val.append(avg_ptl.item())
+
+                cpu = time.time() - cpu
+                print(
+                    "CPU: {}\n"
+                    "Loss: {}\n"
+                    "Average tour length by policy: {}\n"
+                    "Average tour length by rollout: {}\n".format(cpu, loss, avg_ptl, avg_rtl))
+
+                self.log_file.writerow([epoch, losses[-nb_batches:],
+                                        avg_tour_length_batch[-nb_batches:],
+                                        avg_tour_length_epoch[-1],
+                                        avg_ptl.item()
+                                        ])
+
+                if (avg_ptl - avg_rtl) < 0:
+                    # t-test
+                    _, pvalue = ttest_rel(policy_tl, rollout_tl)
+                    pvalue = pvalue / 2  # one-sided ttest [refer to the original implementation]
+                    if pvalue < 0.05:
+                        print("Rollout network update...\n")
+                        self.baseline.load_state_dict(self.model.state_dict())
+                        self.baseline.reset()
+                        print("Generate new validation dataset\n")
+
+                        validation_dataset = VRPDataset(size=self.graph_size, num_samples=self.nb_val_samples)
+
+                        validation_dataloader = DataLoader(validation_dataset, batch_size=self.batch_size, shuffle=True,
+                                                           num_workers=0, pin_memory=True)
+
+            model_name = "RL_{}{}_Epoch_{}.pt".format("vrp", self.graph_size, epoch + 1)
+            torch.save({
+                "epoch": epoch,
+                "model": self.model.state_dict(),
+                "baseline": self.baseline.state_dict(),
+                "optimizer": self.optimizer.state_dict()
+            }, model_name)
+
+        plot.plot_stats(losses, "{}-RL-Losses per batch {}".format("vrp", self.graph_size), "Batch", "Loss")
+        plot.plot_stats(avg_tour_length_epoch,
+                        "{} Average tour length per epoch train {}".format("vrp", self.graph_size),
+                        "Epoch", "Average tour length")
+        plot.plot_stats(avg_tour_length_batch,
+                        "{} Average tour length per batch train {}".format("vrp", self.graph_size),
+                        "Batch", "Average tour length")
+        plot.plot_stats(avg_tl_epoch_val,
+                        "{} Average tour length per epoch validation {}".format("vrp", self.graph_size),
+                        "Epoch", "Average tour length")

src/vrp/data/VRPDataset.py

@@ -0,0 +1,42 @@
+import numpy
+import torch
+from torch.utils.data import Dataset
+
+
+class VRPDataset(Dataset):
+    CAPACITIES = {
+        5: numpy.float32(15),
+        10: numpy.float32(20),
+        20: numpy.float32(30),
+        40: numpy.float32(40),
+        50: numpy.float32(40),
+        100: numpy.float32(50)
+    }
+
+    # code copied from https://github.com/wouterkool/attention-learn-to-route/blob/master/problems/vrp/problem_vrp.py
+    def __init__(self, size, num_samples):
+
+        self.graph_size = size
+        # locations[0] is the depot, locations[1:] are the clients
+        self.locations = []
+
+        # depot demand : demands[0]= 0, client demands : demands[1:] uniform(1,9)
+        self.demands = []
+
+        self.capacities = []
+
+        for sample in range(num_samples):
+            self.capacities.append(self.CAPACITIES[self.graph_size])
+            self.locations.append(torch.FloatTensor(size, 2).uniform_(0, 1))
+            self.demands.append(torch.randint(low=0, high=9, size=(size,), dtype=torch.float32) + 1.0)
+            self.demands[-1][0] = 0
+
+
+        self.size = len(self.locations)
+
+    def __len__(self):
+        return self.size
+
+    def __getitem__(self, idx):
+        return self.locations[idx], self.demands[idx], self.capacities[idx]
+

src/vrp/data/VRPLiteratureDataset.py

@@ -0,0 +1,44 @@
+from typing import Optional
+
+import cvrplib
+import torch
+from torch.utils.data import Dataset
+
+
+class VRPLiteratureDataset(Dataset):
+    def __init__(self,
+                 low: Optional[int] = None,
+                 high: Optional[int] = None):
+        names = cvrplib.list_names(low=low, high=high, vrp_type="cvrp")
+
+        # locations[0] is the depot, locations[1:] are the clients
+        self.locations = []
+
+        # depot demand : demands[0]= 0, client demands : demands[1:] uniform(1,9)
+        self.demands = []
+
+        self.capacities = []
+
+        for name in names:
+            print("Downloading ", name)
+            try:
+                instance = cvrplib.download(name)
+                if instance.coordinates is None:
+                    continue
+                self.locations.append(instance.coordinates)
+                self.demands.append(instance.demands)
+                self.capacities.append(instance.capacity)
+            except:
+                print("Failed to download ", name)
+                continue
+
+        self.size = len(self.locations)
+        assert len(self.locations) == len(self.demands) == len(self.capacities) == self.size
+        self.locations = torch.FloatTensor(self.locations)
+        self.demands = torch.Tensor(self.demands)
+
+    def __len__(self):
+        return self.size
+
+    def __getitem__(self, idx):
+        return self.locations[idx], self.demands[idx], self.capacities[idx]

src/vrp/data/VRPXML100Dataset.py

@@ -0,0 +1,199 @@
+import math
+import random
+
+import torch
+from torch.utils.data import Dataset
+
+
+def distance(x, y):
+    return math.sqrt((x[0] - y[0]) ** 2 + (x[1] - y[1]) ** 2)
+
+
+class VRPXML100Dataset(Dataset):
+    # code copied from the generator of XML100 dataset
+    def __init__(self, size, num_samples, seed=1234):
+
+        self.graph_size = size
+        # locations[0] is the depot, locations[1:] are the clients
+        self.locations = []
+
+        # depot demand : demands[0]= 0, client demands : demands[1:] uniform(1,9)
+        self.demands = []
+
+        self.capacities = []
+
+        # constants
+        max_coord = 1000
+        decay = 40
+        n = size - 1
+        random.seed(seed)
+
+        for _ in range(math.ceil(num_samples / (3 * 3 * 7 * 6))):
+            for rootPos in (1, 2, 3):
+                for custPos in (1, 2, 3):
+                    for demandType in (1, 2, 3, 4, 5, 6, 7):
+                        for avgRouteSize in (1, 2, 3, 4, 5, 6):
+                            n_seeds = random.randint(2, 6)
+                            r_bornes = {1: (3, 5), 2: (5, 8), 3: (8, 12), 4: (12, 16), 5: (16, 25), 6: (25, 50)}
+                            r = random.uniform(r_bornes[avgRouteSize][0], r_bornes[avgRouteSize][1])
+                            coordinates = set()  # set of coordinates for the customers
+
+                            # Root positioning
+                            if rootPos == 1:
+                                x_ = random.randint(0, max_coord)
+                                y_ = random.randint(0, max_coord)
+                            elif rootPos == 2:
+                                x_ = y_ = int(max_coord / 2.0)
+                            elif rootPos == 3:
+                                x_ = y_ = 0
+                            else:
+                                print("Depot Positioning out of range!")
+                                exit(0)
+                            depot = (x_, y_)
+
+                            # Customer positioning
+                            if custPos == 3:
+                                n_rand_cust = int(n / 2.0)
+                            elif custPos == 2:
+                                n_rand_cust = 0
+                            elif custPos == 1:
+                                n_rand_cust = n
+                                n_seeds = 0
+                            else:
+                                print("Costumer Positioning out of range!")
+                                exit(0)
+
+                            n_clust_cust = n - n_rand_cust
+
+                            # Generating random customers
+                            for i in range(1, n_rand_cust + 1):
+                                x_ = random.randint(0, max_coord)
+                                y_ = random.randint(0, max_coord)
+                                while (x_, y_) in coordinates or (x_, y_) == depot:
+                                    x_ = random.randint(0, max_coord)
+                                    y_ = random.randint(0, max_coord)
+                                coordinates.add((x_, y_))
+
+                            n_s = n_rand_cust
+
+                            seeds = []
+                            # Generation of the clustered customers
+                            if n_clust_cust > 0:
+                                if n_clust_cust < n_seeds:
+                                    print("Too many seeds!")
+                                    exit(0)
+
+                                # Generate the seeds
+                                for i in range(n_seeds):
+                                    x_ = random.randint(0, max_coord)
+                                    y_ = random.randint(0, max_coord)
+                                    while (x_, y_) in coordinates or (x_, y_) == depot:
+                                        x_ = random.randint(0, max_coord)
+                                        y_ = random.randint(0, max_coord)
+                                    coordinates.add((x_, y_))
+                                    seeds.append((x_, y_))
+                                n_s = n_s + n_seeds
+
+                                # Determine the seed with maximum sum of weights (w.r.t. all seeds)
+                                max_weight = 0.0
+                                for i, j in seeds:
+                                    w_ij = 0.0
+                                    for i_, j_ in seeds:
+                                        w_ij += 2 ** (- distance((i, j), (i_, j_)) / decay)
+                                    if w_ij > max_weight:
+                                        max_weight = w_ij
+
+                                norm_factor = 1.0 / max_weight
+
+                                # Generate the remaining customers using Accept-reject method
+                                while n_s < n:
+                                    x_ = random.randint(0, max_coord)
+                                    y_ = random.randint(0, max_coord)
+                                    while (x_, y_) in coordinates or (x_, y_) == depot:
+                                        x_ = random.randint(0, max_coord)
+                                        y_ = random.randint(0, max_coord)
+
+                                    weight = 0.0
+                                    for i_, j_ in seeds:
+                                        weight += 2 ** (- distance((x_, y_), (i_, j_)) / decay)
+                                    weight *= norm_factor
+                                    rand = random.uniform(0, 1)
+
+                                    if rand <= weight:  # Will we accept the customer?
+                                        coordinates.add((x_, y_))
+                                        n_s = n_s + 1
+
+                            vertices = [depot] + list(coordinates)  # set of vertices (from now on, the ids are defined)
+
+                            # Demands
+                            demand_min_values = [1, 1, 5, 1, 50, 1, 51, 50, 1]
+                            demand_max_values = [1, 10, 10, 100, 100, 50, 100, 100, 10]
+                            demand_min = demand_min_values[demandType - 1]
+                            demand_max = demand_max_values[demandType - 1]
+                            demand_min_even_quadrant = 51
+                            demand_max_even_quadrant = 100
+                            demand_min_large = 50
+                            demand_max_large = 100
+                            large_per_route = 1.5
+                            demand_min_small = 1
+                            demand_max_small = 10
+
+                            demands = []  # demands
+                            sum_demands = 0
+                            max_demand = 0
+
+                            for i in range(2, n + 2):
+                                j = int((demand_max - demand_min + 1) * random.uniform(0, 1) + demand_min)
+                                if demandType == 6:
+                                    if (vertices[i - 1][0] < max_coord / 2.0 and vertices[i - 1][1] < max_coord / 2.0) or (
+                                            vertices[i - 1][0] >= max_coord / 2.0 and vertices[i - 1][1] >= max_coord / 2.0):
+                                        j = int(
+                                            (demand_max_even_quadrant - demand_min_even_quadrant + 1) * random.uniform(
+                                                0,
+                                                1) + demand_min_even_quadrant)
+                                if demandType == 7:
+                                    if i < (n / r) * large_per_route:
+                                        j = int(
+                                            (demand_max_large - demand_min_large + 1) * random.uniform(0,
+                                                                                                       1) + demand_min_large)
+                                    else:
+                                        j = int(
+                                            (demand_max_small - demand_min_small + 1) * random.uniform(0,
+                                                                                                       1) + demand_min_small)
+                                demands.append(j)
+                                if j > max_demand:
+                                    max_demand = j
+                                sum_demands = sum_demands + j
+
+                            # Generate capacity
+                            if sum_demands == n:
+                                capacity = math.floor(r)
+                            else:
+                                capacity = max(max_demand, math.ceil(r * sum_demands / n))
+
+                            self.capacities.append(capacity)
+
+                            nodes = []
+                            for i, v in enumerate(vertices):
+                                nodes.append((v[0], v[1]))
+                            self.locations.append(nodes)
+
+                            if demandType != 6:
+                                random.shuffle(demands)
+                            demands = [0] + demands
+
+                            self.demands.append(demands)
+
+                            assert len(demands) == len(nodes) == size
+
+                            assert len(self.locations) == len(self.demands) == len(self.capacities)
+
+        self.size = len(self.locations)
+        self.locations = torch.FloatTensor(self.locations)
+        self.demands = torch.Tensor(self.demands)
+
+    def __len__(self):
+        return self.size
+
+    def __getitem__(self, idx):
+        return self.locations[idx], self.demands[idx], self.capacities[idx]

src/vrp/eval.py

@@ -0,0 +1,71 @@
+import os
+
+import torch
+from torch.utils.data import DataLoader
+
+import plot
+import train
+from Model import AttentionModel
+from data.VRPDataset import VRPDataset
+from data.VRPLiteratureDataset import VRPLiteratureDataset
+from data.VRPXML100Dataset import VRPXML100Dataset
+
+if __name__ == "__main__":
+    graph_size = 100
+    batch_size = 100
+    nb_test_samples = 1000
+    test_samples = "xml100"  # "xml100", "random", "literature"
+    n_layers = train.n_layers
+    n_heads = train.n_heads
+    embedding_dim = train.embedding_dim
+    dim_feedforward = train.dim_feedforward
+    decode_mode = "greedy"
+    C = train.C
+    dropout = train.dropout
+    seed = 1234
+    torch.manual_seed(seed)
+
+    print("Generating {} samples of type {} and size {}".format(nb_test_samples, test_samples, graph_size))
+    if test_samples == "xml100":
+        test_dataset = VRPXML100Dataset(size=graph_size, num_samples=nb_test_samples, seed=seed)
+    elif test_samples == "literature":
+        # FIXME : allow to load instances with different sizes
+        test_dataset = VRPLiteratureDataset(low=graph_size - 1, high=graph_size - 1)
+    else:
+        test_dataset = VRPDataset(size=graph_size, num_samples=nb_test_samples)
+
+    print("Number of test samples : ", len(test_dataset))
+
+    test_dataloader = DataLoader(test_dataset, batch_size, shuffle=False, num_workers=1)
+
+    random_model = AttentionModel(embedding_dim, n_layers, n_heads, dim_feedforward, C, dropout)
+
+    tour_lengths = {
+        "random": random_model.test(test_dataloader)
+    }
+
+    datasets = os.listdir("../../pretrained/" + "vrp" + str(graph_size) + "/")
+    for dataset in sorted(datasets, key=lambda x: int(x.split('_')[-1].split('.')[0]) if x.endswith('.pt') else 0):
+        if not dataset.endswith('.pt'): continue
+        print("Testing model : ", dataset)
+        data = torch.load("../../pretrained/vrp" + str(graph_size) + "/" + dataset)
+
+        try:
+            model = AttentionModel(embedding_dim, n_layers, n_heads, dim_feedforward, C, dropout)
+            model.load_state_dict(data["model"])
+            results = model.test(test_dataloader)
+            tour_lengths[int(dataset.split('_')[-1].split('.')[0])] = results
+            print('{} : {} in cpu = {}'.format(dataset, results["avg_tl"], results["cpu"]))
+        except Exception as e:
+            print(e)
+            continue
+
+    plot.plot_stats([results["avg_tl"] for dataset, results in tour_lengths.items()],
+                    "{} Average tour length per epoch evaluation {}".format("vrp", graph_size),
+                    "Epoch", "Average tour length")
+
+    sorted_tour_lengths = sorted(tour_lengths.items(), key=lambda x: x[1]["avg_tl"])
+
+    print('Sorted tour lengths per model')
+    for model_name, results in sorted_tour_lengths:
+        print('{} : {} in cpu = {}'.format(model_name, results["avg_tl"], results["cpu"]))

src/vrp/plot.py

@@ -0,0 +1,12 @@
+import matplotlib.pyplot as plt
+
+
+def plot_stats(stats_array, plot_name, x_name, y_name):
+    plt.plot(stats_array)
+    plt.xlabel(x_name)
+    plt.ylabel(y_name)
+    plt.title(plot_name)
+    plt.savefig(plot_name + ".png")
+    plt.cla()
+    plt.clf()
+

src/vrp/train.py

@@ -0,0 +1,26 @@
+import torch
+
+from Trainer import Trainer
+
+graph_size = 50
+n_epochs = 50
+batch_size = 250
+nb_train_samples = 10000
+nb_val_samples = 1000
+n_layers = 3
+n_heads = 8
+embedding_dim = 128
+dim_feedforward = 512
+C = 10
+dropout = 0.1
+learning_rate = 1e-5
+seed = 1234
+
+if __name__ == "__main__":
+    torch.set_num_threads(4)
+    torch.autograd.set_detect_anomaly(True)
+    torch.manual_seed(seed)
+    trainer = Trainer(graph_size, n_epochs, batch_size, nb_train_samples, nb_val_samples,
+                      n_layers, n_heads, embedding_dim, dim_feedforward, C,
+                      dropout, learning_rate)
+    trainer.train()