Initial commit

farkguidao · Jan 23, 2024 · 6b7c220 · 6b7c220
commit 6b7c220
Show file tree

Hide file tree

Showing 41 changed files with 2,920 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -0,0 +1,87 @@
+# 2D-Ptr
+Source code for paper "2D-Ptr: 2D Array Pointer Network for Solving the Heterogeneous Capacitated Vehicle Routing Problem"
+
+## Dependencies
+
+- Python>=3.8
+- NumPy
+- SciPy
+- [PyTorch](http://pytorch.org/)>=1.12.1
+- tqdm
+- [tensorboard_logger](https://github.com/TeamHG-Memex/tensorboard_logger)
+
+## Quick start
+
+The implementation of the 2D-Ptr model is mainly in the file `./nets/attention_model.py`
+
+For testing HCVRP instances with 60 customers and 5 vehicles (V5-U60) and using pre-trained model:
+
+```shell
+# greedy
+python eval.py data/hcvrp/hcvrp_v5_60_seed24610.pkl --model outputs/hcvrp_v5_60 --obj min-max --decode_strategy greedy --eval_batch_size 1
+# sample1280
+python eval.py data/hcvrp/hcvrp_v5_60_seed24610.pkl --model outputs/hcvrp_v5_60 --obj min-max --decode_strategy sample --width 1280 --eval_batch_size 1
+# sample12800
+python eval.py data/hcvrp/hcvrp_v5_60_seed24610.pkl --model outputs/hcvrp_v5_60 --obj min-max --decode_strategy sample --width 12800 --eval_batch_size 1
+```
+
+Since AAMAS limits the submission file size within 25Mb, we can only provide the pre-trained model on V5-U60 to avoid exceeding the limit.
+
+## Usage
+
+### Generating data
+
+We have provided all the well-generated test datasets in `./data`, and you can also generate each test set by:
+
+```shell
+python generate_data.py --dataset_size 1280 --veh_num 3 --graph_size 40
+```
+
+- The `--graph_size`  and `--veh_num`  represent the number of customers , vehicles and generated instances, respectively.
+
+- The  default random seed is 24610, and you can change it in `./generate_data.py`.
+- The test set will be stored in `./data/hcvrp/`
+
+### Training
+
+For training HCVRP instances with 40 customers and 3 vehicles (V3-U40):
+
+```shell
+python run.py --graph_size 40 --veh_num 3 --baseline rollout --run_name hcvrp_v3_40_rollout --obj min-max
+```
+
+- `--run_name` will be automatically appended with a timestamp, as the unique subpath for logs and checkpoints.
+- The log based on Tensorboard will be stored in `./log/`, and the checkpoint (or the well-trained model) will be stored in `./outputs/`
+- `--obj` represents the objective function, supporting `min-max` and `min-sum`
+
+By default, training will happen on all available GPUs.   Change the code in `./run.py` to only use specific GPUs:
+
+```python
+if __name__ == "__main__":
+    warnings.filterwarnings('ignore')
+    # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    run(get_options())
+```
+
+### Evaluation
+
+you can test a well-trained model on HCVRP instances with any problem size:
+
+```shell
+# greedy
+python eval.py data/hcvrp/hcvrp_v3_40_seed24610.pkl --model outputs/hcvrp_v3_40 --obj min-max --decode_strategy greedy --eval_batch_size 1
+# sample1280
+python eval.py data/hcvrp/hcvrp_v3_40_seed24610.pkl --model outputs/hcvrp_v3_40 --obj min-max --decode_strategy sample --width 1280 --eval_batch_size 1
+# sample12800
+python eval.py data/hcvrp/hcvrp_v3_40_seed24610.pkl --model outputs/hcvrp_v3_40 --obj min-max --decode_strategy sample --width 12800 --eval_batch_size 1
+```
+
+- The `--model`  represents the directory where the used model is located. 
+- The `$filename$.pkl` represents the test set. 
+- The `--width` represents sampling number, which is only available when `--decode_strategy` is `sample`.
+- The `--eval_batch_size` is set to 1 for serial evaluation.
+
+
+
+
+
diff --git a/data/hcvrp/hcvrp_v3_100_seed24610.pkl b/data/hcvrp/hcvrp_v3_100_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v3_40_seed24610.pkl b/data/hcvrp/hcvrp_v3_40_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v3_60_seed24610.pkl b/data/hcvrp/hcvrp_v3_60_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v3_80_seed24610.pkl b/data/hcvrp/hcvrp_v3_80_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v5_100_seed24610.pkl b/data/hcvrp/hcvrp_v5_100_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v5_40_seed24610.pkl b/data/hcvrp/hcvrp_v5_40_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v5_60_seed24610.pkl b/data/hcvrp/hcvrp_v5_60_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v5_80_seed24610.pkl b/data/hcvrp/hcvrp_v5_80_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v7_100_seed24610.pkl b/data/hcvrp/hcvrp_v7_100_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v7_40_seed24610.pkl b/data/hcvrp/hcvrp_v7_40_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v7_60_seed24610.pkl b/data/hcvrp/hcvrp_v7_60_seed24610.pkl
diff --git a/data/hcvrp/hcvrp_v7_80_seed24610.pkl b/data/hcvrp/hcvrp_v7_80_seed24610.pkl
diff --git a/eval.py b/eval.py
@@ -0,0 +1,225 @@
+# used after model is completely trained, and test for results
+
+import math
+import torch
+import os
+import argparse
+import numpy as np
+import itertools
+from tqdm import tqdm
+from utils import load_model, move_to
+from utils.data_utils import save_dataset
+from torch.utils.data import DataLoader
+import time
+from datetime import timedelta
+from utils.functions import parse_softmax_temperature
+import warnings
+
+mp = torch.multiprocessing.get_context('spawn')
+
+
+def get_best(sequences, cost, veh_lists, ids=None, batch_size=None):
+    """
+    Ids contains [0, 0, 0, 1, 1, 2, ..., n, n, n] if 3 solutions found for 0th instance, 2 for 1st, etc
+    :param sequences:
+    :param lengths:
+    :param ids:
+    :return: list with n sequences and list with n lengths of solutions
+    """
+    if ids is None:
+        idx = cost.argmin()
+        return sequences[idx:idx + 1, ...], cost[idx:idx + 1, ...], veh_lists[idx:idx + 1, ...]
+
+    splits = np.hstack([0, np.where(ids[:-1] != ids[1:])[0] + 1])
+    mincosts = np.minimum.reduceat(cost, splits)
+
+    group_lengths = np.diff(np.hstack([splits, len(ids)]))
+    all_argmin = np.flatnonzero(np.repeat(mincosts, group_lengths) == cost)
+    result = np.full(len(group_lengths) if batch_size is None else batch_size, -1, dtype=int)
+
+    result[ids[all_argmin[::-1]]] = all_argmin[::-1]
+
+    return [sequences[i] if i >= 0 else None for i in result], [cost[i] if i >= 0 else math.inf for i in result], [
+        veh_lists[i] if i >= 0 else None for i in result]
+
+
+def eval_dataset_mp(args):
+    (dataset_path, width, softmax_temp, opts, i, num_processes) = args
+
+    model, _ = load_model(opts.model, opts.obj)
+    val_size = opts.val_size // num_processes
+    dataset = model.problem.make_dataset(filename=dataset_path, num_samples=val_size, offset=opts.offset + val_size * i)
+    device = torch.device("cuda:{}".format(i))
+
+    return _eval_dataset(model, dataset, width, softmax_temp, opts, device)
+
+
+def eval_dataset(dataset_path, width, softmax_temp, opts):
+    # Even with multiprocessing, we load the model here since it contains the name where to write results
+    model, _ = load_model(opts.model, opts.obj)
+    use_cuda = torch.cuda.is_available() and not opts.no_cuda
+    if opts.multiprocessing:
+        assert use_cuda, "Can only do multiprocessing with cuda"
+        num_processes = torch.cuda.device_count()
+        assert opts.val_size % num_processes == 0
+
+        with mp.Pool(num_processes) as pool:
+            results = list(itertools.chain.from_iterable(pool.map(
+                eval_dataset_mp,
+                [(dataset_path, width, softmax_temp, opts, i, num_processes) for i in range(num_processes)]
+            )))
+
+    else:
+        device = torch.device("cuda:0" if use_cuda else "cpu")
+        dataset = model.problem.make_dataset(filename=dataset_path, num_samples=opts.val_size, offset=opts.offset)
+        results = _eval_dataset(model, dataset, width, softmax_temp, opts, device)
+
+    # This is parallelism, even if we use multiprocessing (we report as if we did not use multiprocessing, e.g. 1 GPU)
+    parallelism = opts.eval_batch_size
+
+    costs, tours, veh_lists, durations = zip(*results)  # Not really costs since they should be negative
+
+    print("Average cost: {} +- {}".format(np.mean(costs), 2 * np.std(costs) / np.sqrt(len(costs))))
+    print("Average serial duration: {} +- {}".format(
+        np.mean(durations), 2 * np.std(durations) / np.sqrt(len(durations))))
+    print("Average parallel duration: {}".format(np.mean(durations) / parallelism))
+    print("Calculated total duration: {}".format(timedelta(seconds=int(np.sum(durations) / parallelism))))
+    # print('tour is', costs[0], len(tours), len(tours[0]), tours[0])
+    # print('veh', veh_lists[1])
+    # print('tour is', costs[1], len(tours), len(tours[1]), tours[1])
+
+    dataset_basename, ext = os.path.splitext(os.path.split(dataset_path)[-1])
+    model_name = "_".join(os.path.normpath(os.path.splitext(opts.model)[0]).split(os.sep)[-2:])
+    if opts.o is None:
+        results_dir = os.path.join(opts.results_dir, model.problem.NAME, dataset_basename)
+        os.makedirs(results_dir, exist_ok=True)
+
+        out_file = os.path.join(results_dir, "{}-{}-{}{}-t{}-{}-{}{}".format(
+            dataset_basename, model_name,
+            opts.decode_strategy,
+            width if opts.decode_strategy != 'greedy' else '',
+            softmax_temp, opts.offset, opts.offset + len(costs), ext
+        ))
+    else:
+        out_file = opts.o
+
+    assert opts.f or not os.path.isfile(
+        out_file), "File already exists! Try running with -f option to overwrite."
+
+    save_dataset((results, parallelism), out_file)
+
+    return costs, tours, durations
+
+
+def _eval_dataset(model, dataset, width, softmax_temp, opts, device):
+    # print('data', dataset[0])
+    model.to(device)
+    model.eval()
+
+    model.set_decode_type(
+        "greedy" if opts.decode_strategy in ('bs', 'greedy') else "sampling",
+        temp=softmax_temp)
+
+    dataloader = DataLoader(dataset, batch_size=opts.eval_batch_size)
+
+    results = []
+    for batch in tqdm(dataloader, disable=opts.no_progress_bar):
+        batch = move_to(batch, device)
+        start = time.time()
+        with torch.no_grad():
+            if opts.decode_strategy in ('sample', 'greedy'):
+                if opts.decode_strategy == 'greedy':
+                    assert width == 0, "Do not set width when using greedy"
+                    assert opts.eval_batch_size <= opts.max_calc_batch_size, \
+                        "eval_batch_size should be smaller than calc batch size"
+                    batch_rep = 1
+                    iter_rep = 1
+                elif width * opts.eval_batch_size > opts.max_calc_batch_size:
+                    assert opts.eval_batch_size == 1
+                    assert width % opts.max_calc_batch_size == 0
+                    batch_rep = opts.max_calc_batch_size
+                    iter_rep = width // opts.max_calc_batch_size
+                else:
+                    batch_rep = width
+                    iter_rep = 1
+                assert batch_rep > 0
+                # This returns (batch_size, iter_rep shape)
+                sequences, costs, veh_lists = model.sample_many(batch, batch_rep=batch_rep, iter_rep=iter_rep)
+                print('cost', costs)
+                batch_size = len(costs)
+                ids = torch.arange(batch_size, dtype=torch.int64, device=costs.device)
+            else:
+                assert opts.decode_strategy == 'bs'
+
+                cum_log_p, sequences, costs, ids, batch_size = model.beam_search(
+                    batch, beam_size=width,
+                    compress_mask=opts.compress_mask,
+                    max_calc_batch_size=opts.max_calc_batch_size
+                )
+
+        if sequences is None:
+            sequences = [None] * batch_size
+            costs = [math.inf] * batch_size
+            veh_lists = [None] * batch_size
+        else:
+            sequences, costs, veh_lists = get_best(
+                sequences.cpu().numpy(), costs.cpu().numpy(), veh_lists.cpu().numpy(),
+                ids.cpu().numpy() if ids is not None else None,
+                batch_size
+            )
+
+        duration = time.time() - start
+        for seq, cost, veh_list in zip(sequences, costs, veh_lists):
+            if model.problem.NAME in ("hcvrp"):
+                seq = seq.tolist()  # No need to trim as all are same length
+            else:
+                assert False, "Unkown problem: {}".format(model.problem.NAME)
+            # Note VRP only
+            results.append((cost, seq, veh_list, duration))
+
+    return results
+
+
+if __name__ == "__main__":
+    warnings.filterwarnings('ignore')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("datasets", nargs='+', help="Filename of the dataset(s) to evaluate")
+    parser.add_argument("-f", action='store_true', help="Set true to overwrite")
+    parser.add_argument("-o", default=None, help="Name of the results file to write")
+    parser.add_argument('--val_size', type=int, default=10000,
+                        help='Number of instances used for reporting validation performance')
+    parser.add_argument('--offset', type=int, default=0,
+                        help='Offset where to start in dataset (default 0)')
+    parser.add_argument('--eval_batch_size', type=int, default=1024,
+                        help="Batch size to use during (baseline) evaluation")
+    # parser.add_argument('--decode_type', type=str, default='greedy',
+    #                     help='Decode type, greedy or sampling')
+    parser.add_argument('--width', type=int, nargs='+',
+                        help='Sizes of beam to use for beam search (or number of samples for sampling), '
+                             '0 to disable (default), -1 for infinite')
+    parser.add_argument('--decode_strategy', type=str,
+                        help='Beam search (bs), Sampling (sample) or Greedy (greedy)')
+    parser.add_argument('--softmax_temperature', type=parse_softmax_temperature, default=1,
+                        help="Softmax temperature (sampling or bs)")
+    parser.add_argument('--model', type=str)
+    parser.add_argument('--no_cuda', action='store_true', help='Disable CUDA')
+    parser.add_argument('--no_progress_bar', action='store_true', help='Disable progress bar')
+    parser.add_argument('--compress_mask', action='store_true', help='Compress mask into long')
+    parser.add_argument('--max_calc_batch_size', type=int, default=10000000, help='Size for subbatches')
+    parser.add_argument('--results_dir', default='results', help="Name of results directory")
+    parser.add_argument('--obj', default=['min-max', 'min-sum'])
+    parser.add_argument('--multiprocessing', action='store_true',
+                        help='Use multiprocessing to parallelize over multiple GPUs')
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    opts = parser.parse_args()
+
+    assert opts.o is None or (len(opts.datasets) == 1 and len(opts.width) <= 1), \
+        "Cannot specify result filename with more than one dataset or more than one width"
+
+    widths = opts.width if opts.width is not None else [0]
+
+    for width in widths:
+        for dataset_path in opts.datasets:
+            eval_dataset(dataset_path, width, opts.softmax_temperature, opts)
diff --git a/generate_data.py b/generate_data.py
@@ -0,0 +1,53 @@
+import os
+import numpy as np
+from utils.data_utils import check_extension, save_dataset
+import torch
+import pickle
+import argparse
+
+def generate_hcvrp_data(seed,dataset_size, hcvrp_size, veh_num):
+    rnd = np.random.RandomState(seed)
+
+    loc = rnd.uniform(0, 1, size=(dataset_size, hcvrp_size + 1, 2))
+    depot = loc[:, -1]
+    cust = loc[:, :-1]
+    d = rnd.randint(1, 10, [dataset_size, hcvrp_size + 1])
+    d = d[:, :-1]  # the demand of depot is 0, which do not need to generate here
+
+    # vehicle feature
+    speed = rnd.uniform(0.5, 1, size=(dataset_size, veh_num))
+    cap = rnd.randint(20, 41, size=(dataset_size, veh_num))
+
+    data = {
+        'depot': depot.astype(np.float32),
+        'loc': cust.astype(np.float32),
+        'demand': d.astype(np.float32),
+        'capacity': cap.astype(np.float32),
+        'speed': speed.astype(np.float32)
+    }
+    return data
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # parser.add_argument("--filename", help="Filename of the dataset to create (ignores datadir)")
+    parser.add_argument("--dataset_size", type=int, default=1280, help="Size of the dataset")
+    parser.add_argument("--veh_num", type=int, default=3, help="number of the vehicles")
+    parser.add_argument('--graph_size', type=int, default=40,
+                        help="Number of customers")
+
+    opts = parser.parse_args()
+    data_dir = 'data'
+    problem = 'hcvrp'
+    datadir = os.path.join(data_dir, problem)
+    os.makedirs(datadir, exist_ok=True)
+    seed = 24610  # the last seed used for generating HCVRP data
+    # np.random.seed(seed)
+    print(opts.dataset_size, opts.graph_size, opts.veh_num)
+    filename = os.path.join(datadir, '{}_v{}_{}_seed{}.pkl'.format(problem, opts.veh_num, opts.graph_size, seed))
+
+    dataset = generate_hcvrp_data(seed,opts.dataset_size, opts.graph_size, opts.veh_num)
+    print({k:dataset[k][0] for k in dataset})
+    save_dataset(dataset, filename)
+
+
+
diff --git a/main.py b/main.py
diff --git a/nets/__init__.py b/nets/__init__.py