data generation for poor man

henry-yeh · Jul 21, 2024 · 08a6e67 · 08a6e67
1 parent 3163f43
commit 08a6e67
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/eval_atsp/ATSProblemDef.py b/eval_atsp/ATSProblemDef.py
@@ -26,6 +26,7 @@
 """
 
 import torch
+from tqdm import tqdm
 
 
 def get_random_problems(batch_size, node_cnt, problem_gen_params):
@@ -107,11 +108,14 @@ def load_single_problem_from_file(filename, node_cnt, scaler):
     if not os.path.exists("../data/atsp"):
         os.mkdir("../data/atsp")
 
-    torch.manual_seed(1234)   
-    
+    torch.manual_seed(1234)
+
     dataset_size = 30
     for scale in [150, 250, 1000]:
-        problems = get_random_problems(dataset_size, scale, problem_gen_params)
-        print("instances generated")
+        problems = []
+        for inst_id in tqdm(range(dataset_size)):
+            problem = get_random_problems(1, scale, problem_gen_params)
+            problems.append(problem)
+        problems = torch.cat(problems, dim=0)
         torch.save(problems, "../data/atsp/ATSP{}.pt".format(scale))
         print(f"created ../data/atsp/ATSP{scale}.pt")