Update run.py

run.py

@@ -20,16 +20,31 @@ from google_solver.google_model import evaluate_google_model
 with open('params.json', 'r') as f:
     params = json.load(f)
 
-# Save params into a local file for tracking
-with open('params_saved.json', 'w') as f:
-    json.dump(params, f)
-
 # Settings
 device = params['device']
 run_tests = params['run_tests']
 save_results = params['save_results']
 dataset_path = params['dataset_path']
 
+# Create persistent directory
+results_dir = os.path.join('/data', 'results')
+os.makedirs(results_dir, exist_ok=True)
+
+now = datetime.now()
+dt_string = now.strftime("%d-%m-%y %H-%M-%S")
+experiment_path = os.path.join(results_dir, dt_string)
+os.makedirs(experiment_path, exist_ok=True)
+
+train_results_file = os.path.join(experiment_path, 'train_results.txt')
+test_results_file = os.path.join(experiment_path, 'test_results.txt')
+model_path = os.path.join(experiment_path, 'model_state_dict.pt')
+optimizer_path = os.path.join(experiment_path, 'optimizer_state_dict.pt')
+
+with open(train_results_file, 'w') as f: pass
+with open(test_results_file, 'w') as f: pass
+with open(os.path.join(experiment_path, 'params.json'), 'w') as f:
+    json.dump(params, f)
+
 # Dataset sizes
 train_dataset_size = params['train_dataset_size']
 validation_dataset_size = params['validation_dataset_size']
@@ -122,7 +137,7 @@ for epoch in range(params['num_epochs']):
         result = f"{epoch}, {i}, {actor_nn_ratio:.4f}, {actor_baseline_ratio:.4f}, {train_batch_record:.4f}"
         print(result, flush=True)
         if save_results:
-            with open('train_results.txt', 'a') as f:
+            with open(train_results_file, 'a') as f:
                 f.write(result + '\n')
         del batch
 
@@ -144,33 +159,45 @@ for epoch in range(params['num_epochs']):
             baseline_actor.load_state_dict(actor.state_dict())
             print('\nNew baseline record\n')
 
-    # Test every 10 epochs
-    if (epoch % 10 == 0) and run_tests:
-        b = max(int(batch_size // sample_size**2), 1)
-        validation_dataloader = DataLoader(validation_dataset, batch_size=b, collate_fn=validation_dataset.collate)
 
-        tot_cost = 0
-        tot_nn_cost = 0
-        for batch in validation_dataloader:
-            with torch.no_grad():
-                actor.beam_search(sample_size)
-                actor_output = actor(batch)
-                cost = actor_output['total_time']
+# Test block (اختبار النموذج بعد كل 10 epochs)
+if (epoch % 10 == 0) and run_tests:
+    b = max(int(batch_size // sample_size**2), 1)
+    validation_dataloader = DataLoader(validation_dataset, batch_size=b, collate_fn=validation_dataset.collate)
 
-                nn_actor.nearest_neighbors()
-                nn_output = nn_actor(batch)
-                nn_cost = nn_output['total_time']
+    tot_cost = 0
+    tot_nn_cost = 0
+    for batch in validation_dataloader:
+        with torch.no_grad():
+            actor.beam_search(sample_size)
+            actor_output = actor(batch)
+            cost = actor_output['total_time']
 
-            tot_cost += cost.sum().item()
-            tot_nn_cost += nn_cost.sum().item()
+            nn_actor.nearest_neighbors()
+            nn_output = nn_actor(batch)
+            nn_cost = nn_output['total_time']
 
-        ratio = tot_cost / tot_nn_cost
-        validation_record = min(validation_record, ratio)
+        tot_cost += cost.sum().item()
+        tot_nn_cost += nn_cost.sum().item()
 
-        actor_google_ratio = tot_cost / tot_google_scores
-        print(f"\nTest results:\nActor/Google: {actor_google_ratio:.4f}, Actor/NN: {ratio:.4f}, Best NN Ratio: {validation_record:.4f}\n")
+    ratio = tot_cost / tot_nn_cost
+    validation_record = min(validation_record, ratio)
+    actor_google_ratio = tot_cost / tot_google_scores
 
-        if save_results:
-            with open('test_results.txt', 'a') as f:
-                f.write(f"{epoch}, {actor_google_ratio:.4f}, {ratio:.4f}, {validation_record:.4f}\n")
-print("End")
+    print(f"\nTest results:\nActor/Google: {actor_google_ratio:.4f}, Actor/NN: {ratio:.4f}, Best NN Ratio: {validation_record:.4f}\n")
+
+# ✅ حفظ النتائج والنموذج بعد كل اختبار
+if save_results:
+    with open(test_results_file, 'a') as f:
+        f.write(f"{epoch}, {actor_google_ratio:.4f}, {ratio:.4f}, {validation_record:.4f}\n")
+    
+    # حفظ النموذج والـ optimizer دائمًا
+    torch.save(actor.state_dict(), model_path)
+    torch.save(optimizer.state_dict(), optimizer_path)
+
+    # نسخة احتياطية كل 10 epochs
+    torch.save(actor.state_dict(), f"model_epoch_{epoch}.pt")
+    torch.save(optimizer.state_dict(), f"optimizer_epoch_{epoch}.pt")
+
+
+print("End.")