main.py: --resume-phase, --checkpoint-dir, --sinkhorn-batch flags

main.py

@@ -8,11 +8,12 @@ Orchestrates the full experiment pipeline:
   5. Evaluate (W2 for 2D, FID/IS for images)
   6. Visualize results
 
+Supports --resume-phase to continue from a checkpoint after interruption.
+
 Usage:
   python main.py --experiment 2d --dataset 8gaussians --steps 10
-  python main.py --experiment 2d --dataset 8gaussians --steps 10 --device cuda
-  python main.py --experiment 2d --dataset 8gaussians --steps 10 --device cpu
   python main.py --experiment mnist --device cuda
+  python main.py --experiment mnist --resume-phase 2  # skip Phase 1, load checkpoint
   python main.py --experiment cifar10 --device cuda
 
 Reference: arXiv:2401.14069 (Neural Sinkhorn Gradient Flow)
@@ -93,12 +94,10 @@ def run_2d_experiment(config: dict, args):
         data_loader=data_loader,
         config=config,
         device=device,
+        checkpoint_dir=args.checkpoint_dir,
     )
     
-    # Build trajectory pool
     trainer.build_trajectory_pool()
-    
-    # Train velocity field
     history = trainer.train()
     
     train_time = time.time() - start_time
@@ -116,8 +115,6 @@ def run_2d_experiment(config: dict, args):
     )
     
     samples = sampler.sample(num_eval)
-    
-    # Also get trajectory for visualization
     trajectory = sampler.sample_trajectory(min(200, num_eval))
     
     # ---- Evaluation ----
@@ -139,17 +136,14 @@ def run_2d_experiment(config: dict, args):
         title=f"NSGF — {config['dataset']} ({num_steps} steps), W2={metrics.get('w2', 0):.4f}",
         save_path=f"results/nsgf_2d_{config['dataset']}_{num_steps}steps.png",
     )
-    
     plot_2d_trajectory(
         trajectory, test_samples,
         title=f"NSGF Trajectory — {config['dataset']}",
         save_path=f"results/nsgf_trajectory_{config['dataset']}_{num_steps}steps.png",
     )
     
-    # Save model
     torch.save(model.state_dict(), f"results/nsgf_2d_{config['dataset']}.pt")
     logger.info("Model saved.")
-    
     return metrics
 
 
@@ -159,6 +153,8 @@ def run_image_experiment(config: dict, args, dataset_name: str):
     Reference: Section 5.2, Appendix E.2
     """
     device = get_device(args)
+    checkpoint_dir = args.checkpoint_dir
+    resume_phase = args.resume_phase
     logger.info(f"Running {dataset_name.upper()} experiment on {device}")
     
     # Override from args
@@ -168,6 +164,11 @@ def run_image_experiment(config: dict, args, dataset_name: str):
         config["nsgf_training"]["num_iterations"] = args.train_iters
         config["nsf_training"]["num_iterations"] = args.train_iters
         config["time_predictor"]["num_iterations"] = args.train_iters
+    if args.sinkhorn_batch:
+        config["sinkhorn"]["batch_size"] = args.sinkhorn_batch
+    
+    # Inject checkpoint_every into config for trainers to read
+    config["checkpoint_every"] = args.checkpoint_every
     
     # Setup
     data_loader = DatasetLoader(config)
@@ -177,6 +178,25 @@ def run_image_experiment(config: dict, args, dataset_name: str):
     nsf_model = create_velocity_unet(config)
     phase_predictor = create_phase_predictor(config)
     
+    # Load checkpoints if resuming
+    if resume_phase > 1:
+        ckpt_path = os.path.join(checkpoint_dir, f"phase{resume_phase - 1}_complete.pt")
+        if os.path.exists(ckpt_path):
+            ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+            if "nsgf_model_state" in ckpt:
+                nsgf_model.load_state_dict(ckpt["nsgf_model_state"])
+                logger.info(f"Loaded NSGF model from {ckpt_path}")
+            if "nsf_model_state" in ckpt:
+                nsf_model.load_state_dict(ckpt["nsf_model_state"])
+                logger.info(f"Loaded NSF model from {ckpt_path}")
+            if "predictor_state" in ckpt:
+                phase_predictor.load_state_dict(ckpt["predictor_state"])
+                logger.info(f"Loaded phase predictor from {ckpt_path}")
+        else:
+            logger.error(f"Checkpoint not found: {ckpt_path}")
+            logger.error(f"Cannot resume from phase {resume_phase} without phase {resume_phase - 1} checkpoint.")
+            sys.exit(1)
+    
     logger.info(f"NSGF UNet params: {sum(p.numel() for p in nsgf_model.parameters()):,}")
     logger.info(f"NSF UNet params: {sum(p.numel() for p in nsf_model.parameters()):,}")
     logger.info(f"Phase predictor params: {sum(p.numel() for p in phase_predictor.parameters()):,}")
@@ -191,9 +211,10 @@ def run_image_experiment(config: dict, args, dataset_name: str):
         data_loader=data_loader,
         config=config,
         device=device,
+        checkpoint_dir=checkpoint_dir,
     )
     
-    results = pp_trainer.train_all()
+    results = pp_trainer.train_all(resume_phase=resume_phase)
     
     train_time = time.time() - start_time
     logger.info(f"Training completed in {train_time:.1f}s")
@@ -215,7 +236,6 @@ def run_image_experiment(config: dict, args, dataset_name: str):
     )
     
     logger.info(f"Generating {num_gen} samples...")
-    # Generate in batches to avoid OOM
     batch_size = 128
     all_samples = []
     for i in range(0, num_gen, batch_size):
@@ -225,10 +245,7 @@ def run_image_experiment(config: dict, args, dataset_name: str):
     generated = torch.cat(all_samples, dim=0)
     
     # ---- Evaluation ----
-    # Get test set
-    eval_loader = data_loader
-    test_samples = eval_loader.get_test_samples(num_gen, device="cpu")
-    
+    test_samples = data_loader.get_test_samples(num_gen, device="cpu")
     evaluator = Evaluation(config, device)
     metrics = evaluator.evaluate(generated, test_samples)
     
@@ -248,7 +265,7 @@ def run_image_experiment(config: dict, args, dataset_name: str):
         save_path=f"results/nsgf_pp_{dataset_name}_samples.png",
     )
     
-    # Save models
+    # Save final models
     torch.save(nsgf_model.state_dict(), f"results/nsgf_{dataset_name}_nsgf.pt")
     torch.save(nsf_model.state_dict(), f"results/nsgf_{dataset_name}_nsf.pt")
     torch.save(phase_predictor.state_dict(), f"results/nsgf_{dataset_name}_predictor.pt")
@@ -268,11 +285,20 @@ def main():
     parser.add_argument("--steps", type=int, default=None, help="Number of flow steps")
     parser.add_argument("--pool-batches", type=int, default=None, help="Pool building batches")
     parser.add_argument("--train-iters", type=int, default=None, help="Training iterations")
+    parser.add_argument("--sinkhorn-batch", type=int, default=None,
+                        help="Sinkhorn batch size for pool building (reduce for OOM)")
     parser.add_argument("--config", type=str, default="config.yaml", help="Config file path")
     parser.add_argument("--seed", type=int, default=42, help="Random seed")
     parser.add_argument("--device", type=str, default=None,
                         choices=["cpu", "cuda"],
                         help="Force device (default: auto-detect)")
+    parser.add_argument("--checkpoint-dir", type=str, default="checkpoints",
+                        help="Directory for saving/loading checkpoints")
+    parser.add_argument("--checkpoint-every", type=int, default=5000,
+                        help="Save checkpoint every N training steps")
+    parser.add_argument("--resume-phase", type=int, default=1,
+                        choices=[1, 2, 3],
+                        help="Resume from phase N (loads phase N-1 checkpoint)")
     
     args = parser.parse_args()