update code

.gitignore

@@ -20,3 +20,4 @@ pretrained/
 utils_plot*
 imgs/
 1.md
+

POMO/TSP/TSPTrainer.py

@@ -77,6 +77,9 @@ def run(self):
             self.scheduler.step()  # LR Decay
             self.result_log.append('train_score', epoch, train_score)
             self.result_log.append('train_loss', epoch, train_loss)
+            # Val
+            no_aug_score = self._fast_val(self.model, val_episodes=64)
+            self.result_log.append('val_score', epoch, no_aug_score)
 
             # Logs & Checkpoint
             elapsed_time_str, remain_time_str = self.time_estimator.get_est_string(epoch, self.trainer_params['epochs'])
@@ -96,10 +99,6 @@ def run(self):
                                     self.result_log, labels=['train_loss'])
 
             if all_done or (epoch % model_save_interval) == 0:
-                # val
-                val_episodes = 64
-                no_aug_score = self._fast_val(self.model, val_episodes=val_episodes)
-                print(">> validation results: {} over {} instances".format(no_aug_score, val_episodes))
                 # save checkpoint
                 self.logger.info("Saving trained_model")
                 checkpoint_dict = {

POMO/TSP/TSPTrainer_Meta.py

@@ -1,5 +1,6 @@
 import copy
 import math
+import time
 import random
 import torch
 from logging import getLogger
@@ -55,9 +56,9 @@ def __init__(self,
         self.meta_model = Model(**self.model_params)
         self.alpha = self.meta_params['alpha']  # for reptile
         self.task_set = generate_task_set(self.meta_params)
-        assert self.trainer_params['meta_params']['epochs'] == math.ceil((self.trainer_params['epochs'] * self.trainer_params['train_episodes']) / (
-                    self.trainer_params['meta_params']['B'] * self.trainer_params['meta_params']['k'] *
-                    self.trainer_params['meta_params']['meta_batch_size'])), ">> meta-learning iteration does not match with POMO!"
+        # assert self.trainer_params['meta_params']['epochs'] == math.ceil((self.trainer_params['epochs'] * self.trainer_params['train_episodes']) / (
+        #             self.trainer_params['meta_params']['B'] * self.trainer_params['meta_params']['k'] *
+        #             self.trainer_params['meta_params']['meta_batch_size'])), ">> meta-learning iteration does not match with POMO!"
 
         # Restore
         self.start_epoch = 1
@@ -77,6 +78,7 @@ def __init__(self,
 
     def run(self):
 
+        start_time = time.time()
         self.time_estimator.reset(self.start_epoch)
         for epoch in range(self.start_epoch, self.meta_params['epochs']+1):
             self.logger.info('=================================================================')
@@ -94,7 +96,10 @@ def run(self):
             self.logger.info("Epoch {:3d}/{:3d}({:.2f}%): Time Est.: Elapsed[{}], Remain[{}], Val Score: {:.4f}".format(
                 epoch, self.meta_params['epochs'], epoch/self.meta_params['epochs']*100, elapsed_time_str, remain_time_str, no_aug_score))
 
-            all_done = (epoch == self.meta_params['epochs'])
+            if self.trainer_params['stop_criterion'] == "epochs":
+                all_done = (epoch == self.meta_params['epochs'])
+            else:
+                all_done = (time.time() - start_time) >= self.trainer_params['time_limit']
             model_save_interval = self.trainer_params['logging']['model_save_interval']
             img_save_interval = self.trainer_params['logging']['img_save_interval']
 

POMO/TSP/TSPTrainer_pomo.py

@@ -1,5 +1,6 @@
 import copy
 import math
+import time
 import random
 import torch
 from logging import getLogger
@@ -53,9 +54,9 @@ def __init__(self,
         self.meta_model = Model(**self.model_params)
         self.optimizer = Optimizer(self.meta_model.parameters(), **self.optimizer_params['optimizer'])
         self.task_set = generate_task_set(self.meta_params)
-        assert self.trainer_params['meta_params']['epochs'] == math.ceil((self.trainer_params['epochs'] * self.trainer_params['train_episodes']) / (
-                    self.trainer_params['meta_params']['B'] * self.trainer_params['meta_params']['k'] *
-                    self.trainer_params['meta_params']['meta_batch_size'])), ">> meta-learning iteration does not match with POMO!"
+        # assert self.trainer_params['meta_params']['epochs'] == math.ceil((self.trainer_params['epochs'] * self.trainer_params['train_episodes']) / (
+        #             self.trainer_params['meta_params']['B'] * self.trainer_params['meta_params']['k'] *
+        #             self.trainer_params['meta_params']['meta_batch_size'])), ">> meta-learning iteration does not match with POMO!"
 
         # Restore
         self.start_epoch = 1
@@ -75,6 +76,7 @@ def __init__(self,
 
     def run(self):
 
+        start_time = time.time()
         self.time_estimator.reset(self.start_epoch)
         for epoch in range(self.start_epoch, self.meta_params['epochs']+1):
             self.logger.info('=================================================================')
@@ -92,7 +94,10 @@ def run(self):
             self.logger.info("Epoch {:3d}/{:3d}({:.2f}%): Time Est.: Elapsed[{}], Remain[{}], Val Score: {:.4f}".format(
                 epoch, self.meta_params['epochs'], epoch/self.meta_params['epochs']*100, elapsed_time_str, remain_time_str, no_aug_score))
 
-            all_done = (epoch == self.meta_params['epochs'])
+            if self.trainer_params['stop_criterion'] == "epochs":
+                all_done = (epoch == self.meta_params['epochs'])
+            else:
+                all_done = (time.time() - start_time) >= self.trainer_params['time_limit']
             model_save_interval = self.trainer_params['logging']['model_save_interval']
             img_save_interval = self.trainer_params['logging']['img_save_interval']
 

POMO/TSP/TSProblemDef.py

@@ -229,10 +229,15 @@ def generate_tsp_dist(n_samples, n_nodes, distribution):
         os.makedirs(path)
     seed_everything(seed=2023)
 
+    # var-dist test data
     for dist in ["uniform", "uniform_rectangle", "gaussian", "cluster", "diagonal", "tsplib"]:
         print(">> Generating TSP instances following {} distribution!".format(dist))
         get_random_problems(20000, 100, distribution=dist, path=path)
 
+    for s in [150, 200]:
+        print(">> Generating TSP instances of size {}!".format(s))
+        get_random_problems(20000, s, distribution="uniform", path=path)
+
     # data = generate_gaussian_mixture_tsp(dataset_size=64, graph_size=100, num_modes=1, cdist=1)
     # data = load_dataset("../../data/TSP/tsp100_cluster.pkl")
     # print(type(data), data.size(), data)

POMO/TSP/train_n100.py

@@ -48,7 +48,9 @@
     'use_cuda': USE_CUDA,
     'cuda_device_num': CUDA_DEVICE_NUM,
     'seed': 1234,
-    'epochs': 500,  # will be overridden if meta_params['enable'] is True
+    'epochs': 500,
+    'time_limit': 86400,
+    'stop_criterion': 'epochs',  # epochs or time
     'train_episodes': 100000,  # number of instances per epoch
     'train_batch_size': 64,
     'logging': {