Test Model on TSPLib¶
In this notebook, we will test the trained model's performance on the TSPLib benchmark. We will use the trained model from the previous notebook.
TSPLib is a library of sample instances for the TSP (and related problems) from various sources and of various types. In the TSPLib, there are several problems, including TSP, HCP, ATSP, etc. In this notebook, we will focus on testing the model on the TSP problem.
Before we start¶
Before we test the model on TSPLib dataset, we need to prepare the dataset first by the following steps:
Step 1. You may come to here to download the symmetric traveling salesman problem data in TSPLib dataset and unzip to a folder;
Note that the downloaded data is gzip
file with the following file tree:
.
└── ALL_tsp/
├── a280.opt.tour.gz
├── a280.tsp.gz
├── ali535.tsp.gz
└── ... (other problems)
We need to unzip the gzip
file to get the .tsp
and .opt.tour
files. We can use the following command to unzip them to the same folder:
find . -name "*.gz" -exec gunzip {} +
After doing this, we will get the following file tree:
.
└── ALL_tsp/
├── a280.opt.tour
├── a280.tsp
├── ali535.tsp
└── ... (other problems)
Step 2. To read the TSPLib problem and optimal solution, we choose to use the tsplib95
package. Use pip install tsplib95
to install the package.
Installation¶
Uncomment the following line to install the package from PyPI. Remember to choose a GPU runtime for faster training!
Note: You may need to restart the runtime in Colab after this
# !pip install rl4co[graph] # include torch-geometric
## NOTE: to install latest version from Github (may be unstable) install from source instead:
# !pip install git+https://github.com/ai4co/rl4co.git
# Install the `tsplib95` package
# !pip install tsplib95
Imports¶
%load_ext autoreload
%autoreload 2
import os
import re
import torch
from rl4co.envs import TSPEnv, CVRPEnv
from rl4co.models.zoo.am import AttentionModel
from rl4co.utils.trainer import RL4COTrainer
from rl4co.utils.decoding import get_log_likelihood
from rl4co.models.zoo import EAS, EASLay, EASEmb, ActiveSearch
from math import ceil
from einops import repeat
from tsplib95.loaders import load_problem, load_solution
/home/cbhua/.local/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
Load a trained model¶
# Load from checkpoint; alternatively, simply instantiate a new model
# Note the model is trained for TSP problem
checkpoint_path = "../tsp-20.ckpt" # modify the path to your checkpoint file
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# load checkpoint
# checkpoint = torch.load(checkpoint_path)
lit_model = AttentionModel.load_from_checkpoint(checkpoint_path, load_baseline=False)
policy, env = lit_model.policy, lit_model.env
policy = policy.to(device)
/home/cbhua/miniconda3/envs/rl4co-user/lib/python3.10/site-packages/lightning/pytorch/utilities/parsing.py:198: Attribute 'env' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['env'])`. /home/cbhua/miniconda3/envs/rl4co-user/lib/python3.10/site-packages/lightning/pytorch/utilities/parsing.py:198: Attribute 'policy' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['policy'])`. /home/cbhua/miniconda3/envs/rl4co-user/lib/python3.10/site-packages/lightning/pytorch/core/saving.py:177: Found keys that are not in the model state dict but in the checkpoint: ['baseline.baseline.model.encoder.init_embedding.init_embed.weight', 'baseline.baseline.model.encoder.init_embedding.init_embed.bias', 'baseline.baseline.model.encoder.net.layers.0.0.module.Wqkv.weight', 'baseline.baseline.model.encoder.net.layers.0.0.module.Wqkv.bias', 'baseline.baseline.model.encoder.net.layers.0.0.module.out_proj.weight', 'baseline.baseline.model.encoder.net.layers.0.0.module.out_proj.bias', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.0.2.module.0.weight', 'baseline.baseline.model.encoder.net.layers.0.2.module.0.bias', 'baseline.baseline.model.encoder.net.layers.0.2.module.2.weight', 'baseline.baseline.model.encoder.net.layers.0.2.module.2.bias', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.1.0.module.Wqkv.weight', 'baseline.baseline.model.encoder.net.layers.1.0.module.Wqkv.bias', 'baseline.baseline.model.encoder.net.layers.1.0.module.out_proj.weight', 'baseline.baseline.model.encoder.net.layers.1.0.module.out_proj.bias', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.1.2.module.0.weight', 'baseline.baseline.model.encoder.net.layers.1.2.module.0.bias', 'baseline.baseline.model.encoder.net.layers.1.2.module.2.weight', 'baseline.baseline.model.encoder.net.layers.1.2.module.2.bias', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.2.0.module.Wqkv.weight', 'baseline.baseline.model.encoder.net.layers.2.0.module.Wqkv.bias', 'baseline.baseline.model.encoder.net.layers.2.0.module.out_proj.weight', 'baseline.baseline.model.encoder.net.layers.2.0.module.out_proj.bias', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.2.2.module.0.weight', 'baseline.baseline.model.encoder.net.layers.2.2.module.0.bias', 'baseline.baseline.model.encoder.net.layers.2.2.module.2.weight', 'baseline.baseline.model.encoder.net.layers.2.2.module.2.bias', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.num_batches_tracked', 'baseline.baseline.model.decoder.context_embedding.W_placeholder', 'baseline.baseline.model.decoder.context_embedding.project_context.weight', 'baseline.baseline.model.decoder.project_node_embeddings.weight', 'baseline.baseline.model.decoder.project_fixed_context.weight', 'baseline.baseline.model.decoder.logit_attention.project_out.weight']
Load tsp problems¶
Note that in the TSPLib, only part of the problems have optimal solutions with the same problem name but with .opt.tour
suffix. For example, a280.tsp
has the optimal solution a280.opt.tour
.
# Load the problem from TSPLib
tsplib_dir = './tsplib'# modify this to the directory of your prepared files
files = os.listdir(tsplib_dir)
problem_files_full = [file for file in files if file.endswith('.tsp')]
# Load the optimal solution files from TSPLib
solution_files = [file for file in files if file.endswith('.opt.tour')]
# Utils function
def normalize_coord(coord:torch.Tensor) -> torch.Tensor:
x, y = coord[:, 0], coord[:, 1]
x_min, x_max = x.min(), x.max()
y_min, y_max = y.min(), y.max()
x_scaled = (x - x_min) / (x_max - x_min)
y_scaled = (y - y_min) / (y_max - y_min)
coord_scaled = torch.stack([x_scaled, y_scaled], dim=1)
return coord_scaled
Test the greedy¶
Note that run all experiments will take long time and require large VRAM. For simple testing, we can use a subset of the problems.
# Customized problem cases
problem_files_custom = [
"eil51.tsp", "berlin52.tsp", "st70.tsp", "eil76.tsp",
"pr76.tsp", "rat99.tsp", "kroA100.tsp", "kroB100.tsp",
"kroC100.tsp", "kroD100.tsp", "kroE100.tsp", "rd100.tsp",
"eil101.tsp", "lin105.tsp", "pr124.tsp", "bier127.tsp",
"ch130.tsp", "pr136.tsp", "pr144.tsp", "kroA150.tsp",
"kroB150.tsp", "pr152.tsp", "u159.tsp", "rat195.tsp",
"kroA200.tsp", "ts225.tsp", "tsp225.tsp", "pr226.tsp"
]
# problem_files = problem_files_full # if you want to test on all the problems
problem_files = problem_files_custom # if you want to test on the customized problems
for problem_idx in range(len(problem_files)):
problem = load_problem(os.path.join(tsplib_dir, problem_files[problem_idx]))
# NOTE: in some problem files (e.g. hk48), the node coordinates are not available
# we temporarily skip these problems
if not len(problem.node_coords):
continue
# Load the node coordinates
coords = []
for _, v in problem.node_coords.items():
coords.append(v)
coords = torch.tensor(coords).float().to(device) # [n, 2]
coords_norm = normalize_coord(coords)
# Prepare the tensordict
batch_size = 2
td = env.reset(batch_size=(batch_size,)).to(device)
td['locs'] = repeat(coords_norm, 'n d -> b n d', b=batch_size, d=2)
td['action_mask'] = torch.ones(batch_size, coords_norm.shape[0], dtype=torch.bool)
# Get the solution from the policy
with torch.no_grad():
out = policy(
td.clone(),
decode_type="greedy",
return_actions=True,
num_starts=0
)
# Calculate the cost on the original scale
td['locs'] = repeat(coords, 'n d -> b n d', b=batch_size, d=2)
neg_reward = env.get_reward(td, out['actions'])
cost = ceil(-1 * neg_reward[0].item())
# Check if there exists an optimal solution
try:
# Load the optimal solution
solution = load_solution(os.path.join(tsplib_dir, problem_files[problem_idx][:-4] + '.opt.tour'))
matches = re.findall(r'\((.*?)\)', solution.comment)
# NOTE: in some problem solution file (e.g. ch130), the optimal cost is not writen with a brace
# we temporarily skip to calculate the gap for these problems
optimal_cost = int(matches[0])
gap = (cost - optimal_cost) / optimal_cost
print(f'Problem: {problem_files[problem_idx][:-4]:<10} Cost: {cost:<10} Optimal Cost: {optimal_cost:<10}\t Gap: {gap:.2%}')
except:
continue
finally:
print(f'problem: {problem_files[problem_idx][:-4]:<10} cost: {cost:<10}')
/tmp/ipykernel_3883036/2632546508.py:5: DeprecationWarning: Call to deprecated function (or staticmethod) load_problem. (Will be removed in newer versions. Use `tsplib95.load` instead.) -- Deprecated since version 7.0.0. problem = load_problem(os.path.join(tsplib_dir, problem_files[problem_idx])) /tmp/ipykernel_3883036/2632546508.py:43: DeprecationWarning: Call to deprecated function (or staticmethod) load_solution. (Will be removed in newer versions. Use `tsplib95.load` instead.) -- Deprecated since version 7.0.0. solution = load_solution(os.path.join(tsplib_dir, problem_files[problem_idx][:-4] + '.opt.tour'))
Problem: eil51 Cost: 493 Optimal Cost: 426 Gap: 15.73% problem: eil51 cost: 493 problem: berlin52 cost: 8957 Problem: st70 Cost: 806 Optimal Cost: 675 Gap: 19.41% problem: st70 cost: 806 Problem: eil76 Cost: 693 Optimal Cost: 538 Gap: 28.81% problem: eil76 cost: 693 Problem: pr76 Cost: 132292 Optimal Cost: 108159 Gap: 22.31% problem: pr76 cost: 132292 problem: rat99 cost: 2053 Problem: kroA100 Cost: 30791 Optimal Cost: 21282 Gap: 44.68% problem: kroA100 cost: 30791 problem: kroB100 cost: 30347 Problem: kroC100 Cost: 28339 Optimal Cost: 20749 Gap: 36.58% problem: kroC100 cost: 28339 Problem: kroD100 Cost: 27600 Optimal Cost: 21294 Gap: 29.61% problem: kroD100 cost: 27600 problem: kroE100 cost: 28396 Problem: rd100 Cost: 10695 Optimal Cost: 7910 Gap: 35.21% problem: rd100 cost: 10695 problem: eil101 cost: 919 Problem: lin105 Cost: 21796 Optimal Cost: 14379 Gap: 51.58% problem: lin105 cost: 21796 problem: pr124 cost: 75310 problem: bier127 cost: 177471 problem: ch130 cost: 8169 problem: pr136 cost: 135974 problem: pr144 cost: 71599 problem: kroA150 cost: 40376 problem: kroB150 cost: 37076 problem: pr152 cost: 94805 problem: u159 cost: 64768 problem: rat195 cost: 4465 problem: kroA200 cost: 44181 problem: ts225 cost: 210475 Problem: tsp225 Cost: 6212 Optimal Cost: 3919 Gap: 58.51% problem: tsp225 cost: 6212 problem: pr226 cost: 98849
Test the Augmentation¶
# problem_files = problem_files_full # if you want to test on all the problems
problem_files = problem_files_custom # if you want to test on the customized problems
# Import augmented utils
from rl4co.data.transforms import (
StateAugmentation as SymmetricStateAugmentation)
from rl4co.utils.ops import batchify, unbatchify
num_augment = 100
augmentation = SymmetricStateAugmentation(num_augment=num_augment)
for problem_idx in range(len(problem_files)):
problem = load_problem(os.path.join(tsplib_dir, problem_files[problem_idx]))
# NOTE: in some problem files (e.g. hk48), the node coordinates are not available
# we temporarily skip these problems
if not len(problem.node_coords):
continue
# Load the node coordinates
coords = []
for _, v in problem.node_coords.items():
coords.append(v)
coords = torch.tensor(coords).float().to(device) # [n, 2]
coords_norm = normalize_coord(coords)
# Prepare the tensordict
batch_size = 2
td = env.reset(batch_size=(batch_size,)).to(device)
td['locs'] = repeat(coords_norm, 'n d -> b n d', b=batch_size, d=2)
td['action_mask'] = torch.ones(batch_size, coords_norm.shape[0], dtype=torch.bool)
# Augmentation
td = augmentation(td)
# Get the solution from the policy
with torch.no_grad():
out = policy(
td.clone(),
decode_type="greedy",
return_actions=True,
num_starts=0
)
# Calculate the cost on the original scale
coords_repeat = repeat(coords, 'n d -> b n d', b=batch_size, d=2)
td['locs'] = batchify(coords_repeat, num_augment)
reward = env.get_reward(td, out['actions'])
reward = unbatchify(reward, num_augment)
cost = ceil(-1 * torch.max(reward).item())
# Check if there exists an optimal solution
try:
# Load the optimal solution
solution = load_solution(os.path.join(tsplib_dir, problem_files[problem_idx][:-4] + '.opt.tour'))
matches = re.findall(r'\((.*?)\)', solution.comment)
# NOTE: in some problem solution file (e.g. ch130), the optimal cost is not writen with a brace
# we temporarily skip to calculate the gap for these problems
optimal_cost = int(matches[0])
gap = (cost - optimal_cost) / optimal_cost
print(f'Problem: {problem_files[problem_idx][:-4]}\t Cost: {cost}\t Optimal Cost: {optimal_cost}\t Gap: {gap:.2%}')
except:
continue
finally:
print(f'problem: {problem_files[problem_idx][:-4]}\t cost: {cost}\t')
/tmp/ipykernel_3883036/2898406631.py:13: DeprecationWarning: Call to deprecated function (or staticmethod) load_problem. (Will be removed in newer versions. Use `tsplib95.load` instead.) -- Deprecated since version 7.0.0. problem = load_problem(os.path.join(tsplib_dir, problem_files[problem_idx])) /tmp/ipykernel_3883036/2898406631.py:56: DeprecationWarning: Call to deprecated function (or staticmethod) load_solution. (Will be removed in newer versions. Use `tsplib95.load` instead.) -- Deprecated since version 7.0.0. solution = load_solution(os.path.join(tsplib_dir, problem_files[problem_idx][:-4] + '.opt.tour'))
Problem: eil51 Cost: 457 Optimal Cost: 426 Gap: 7.28% problem: eil51 cost: 457 problem: berlin52 cost: 8256 Problem: st70 Cost: 777 Optimal Cost: 675 Gap: 15.11% problem: st70 cost: 777 Problem: eil76 Cost: 652 Optimal Cost: 538 Gap: 21.19% problem: eil76 cost: 652 Problem: pr76 Cost: 124939 Optimal Cost: 108159 Gap: 15.51% problem: pr76 cost: 124939 problem: rat99 cost: 1614 Problem: kroA100 Cost: 27694 Optimal Cost: 21282 Gap: 30.13% problem: kroA100 cost: 27694 problem: kroB100 cost: 28244 Problem: kroC100 Cost: 25032 Optimal Cost: 20749 Gap: 20.64% problem: kroC100 cost: 25032 Problem: kroD100 Cost: 26811 Optimal Cost: 21294 Gap: 25.91% problem: kroD100 cost: 26811 problem: kroE100 cost: 27831 Problem: rd100 Cost: 9657 Optimal Cost: 7910 Gap: 22.09% problem: rd100 cost: 9657 problem: eil101 cost: 781 Problem: lin105 Cost: 18769 Optimal Cost: 14379 Gap: 30.53% problem: lin105 cost: 18769 problem: pr124 cost: 72115 problem: bier127 cost: 154518 problem: ch130 cost: 7543 problem: pr136 cost: 128134 problem: pr144 cost: 69755 problem: kroA150 cost: 35967 problem: kroB150 cost: 35196 problem: pr152 cost: 88602 problem: u159 cost: 59484 problem: rat195 cost: 3631 problem: kroA200 cost: 42061 problem: ts225 cost: 196545 Problem: tsp225 Cost: 5680 Optimal Cost: 3919 Gap: 44.93% problem: tsp225 cost: 5680 problem: pr226 cost: 94540
Test the Sampling¶
# problem_files = problem_files_full # if you want to test on all the problems
problem_files = problem_files_custom # if you want to test on the customized problems
# Parameters for sampling
num_samples = 100
softmax_temp = 0.05
for problem_idx in range(len(problem_files)):
problem = load_problem(os.path.join(tsplib_dir, problem_files[problem_idx]))
# NOTE: in some problem files (e.g. hk48), the node coordinates are not available
# we temporarily skip these problems
if not len(problem.node_coords):
continue
# Load the node coordinates
coords = []
for _, v in problem.node_coords.items():
coords.append(v)
coords = torch.tensor(coords).float().to(device) # [n, 2]
coords_norm = normalize_coord(coords)
# Prepare the tensordict
batch_size = 2
td = env.reset(batch_size=(batch_size,)).to(device)
td['locs'] = repeat(coords_norm, 'n d -> b n d', b=batch_size, d=2)
td['action_mask'] = torch.ones(batch_size, coords_norm.shape[0], dtype=torch.bool)
# Sampling
td = batchify(td, num_samples)
# Get the solution from the policy
with torch.no_grad():
out = policy(
td.clone(),
decode_type="sampling",
return_actions=True,
num_starts=0,
softmax_temp=softmax_temp
)
# Calculate the cost on the original scale
coords_repeat = repeat(coords, 'n d -> b n d', b=batch_size, d=2)
td['locs'] = batchify(coords_repeat, num_samples)
reward = env.get_reward(td, out['actions'])
reward = unbatchify(reward, num_samples)
cost = ceil(-1 * torch.max(reward).item())
# Check if there exists an optimal solution
try:
# Load the optimal solution
solution = load_solution(os.path.join(tsplib_dir, problem_files[problem_idx][:-4] + '.opt.tour'))
matches = re.findall(r'\((.*?)\)', solution.comment)
# NOTE: in some problem solution file (e.g. ch130), the optimal cost is not writen with a brace
# we temporarily skip to calculate the gap for these problems
optimal_cost = int(matches[0])
gap = (cost - optimal_cost) / optimal_cost
print(f'Problem: {problem_files[problem_idx][:-4]}\t Cost: {cost}\t Optimal Cost: {optimal_cost}\t Gap: {gap:.2%}')
except:
continue
finally:
print(f'problem: {problem_files[problem_idx][:-4]}\t cost: {cost}\t')
/tmp/ipykernel_3883036/2154301274.py:9: DeprecationWarning: Call to deprecated function (or staticmethod) load_problem. (Will be removed in newer versions. Use `tsplib95.load` instead.) -- Deprecated since version 7.0.0. problem = load_problem(os.path.join(tsplib_dir, problem_files[problem_idx])) /tmp/ipykernel_3883036/2154301274.py:53: DeprecationWarning: Call to deprecated function (or staticmethod) load_solution. (Will be removed in newer versions. Use `tsplib95.load` instead.) -- Deprecated since version 7.0.0. solution = load_solution(os.path.join(tsplib_dir, problem_files[problem_idx][:-4] + '.opt.tour'))
Problem: eil51 Cost: 482 Optimal Cost: 426 Gap: 13.15% problem: eil51 cost: 482 problem: berlin52 cost: 8955 Problem: st70 Cost: 794 Optimal Cost: 675 Gap: 17.63% problem: st70 cost: 794 Problem: eil76 Cost: 673 Optimal Cost: 538 Gap: 25.09% problem: eil76 cost: 673 Problem: pr76 Cost: 127046 Optimal Cost: 108159 Gap: 17.46% problem: pr76 cost: 127046 problem: rat99 cost: 1886 Problem: kroA100 Cost: 29517 Optimal Cost: 21282 Gap: 38.69% problem: kroA100 cost: 29517 problem: kroB100 cost: 28892 Problem: kroC100 Cost: 26697 Optimal Cost: 20749 Gap: 28.67% problem: kroC100 cost: 26697 Problem: kroD100 Cost: 27122 Optimal Cost: 21294 Gap: 27.37% problem: kroD100 cost: 27122 problem: kroE100 cost: 28016 Problem: rd100 Cost: 10424 Optimal Cost: 7910 Gap: 31.78% problem: rd100 cost: 10424 problem: eil101 cost: 837 Problem: lin105 Cost: 19618 Optimal Cost: 14379 Gap: 36.44% problem: lin105 cost: 19618 problem: pr124 cost: 74699 problem: bier127 cost: 170255 problem: ch130 cost: 7985 problem: pr136 cost: 129964 problem: pr144 cost: 70477 problem: kroA150 cost: 37185 problem: kroB150 cost: 35172 problem: pr152 cost: 97244 problem: u159 cost: 59792 problem: rat195 cost: 4325 problem: kroA200 cost: 42059 problem: ts225 cost: 205982 Problem: tsp225 Cost: 5970 Optimal Cost: 3919 Gap: 52.33% problem: tsp225 cost: 5970 problem: pr226 cost: 103135