Upload main.py

main.py

@@ -1,10 +1,19 @@
 """main.py — Entry point for NSGF/NSGF++ experiments.
 
+Orchestrates the full experiment pipeline:
+  1. Load configuration
+  2. Set up dataset, model, trainer
+  3. Train (build pool → velocity matching → [NSF → phase predictor for NSGF++])
+  4. Generate samples
+  5. Evaluate (W2 for 2D, FID/IS for images)
+  6. Visualize results
+
 Usage:
   python main.py --experiment 2d --dataset 8gaussians --steps 10
-  python main.py --experiment 2d --dataset moons --steps 100
-  python main.py --experiment mnist
-  python main.py --experiment cifar10
+  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 cifar10 --device cuda
 
 Reference: arXiv:2401.14069 (Neural Sinkhorn Gradient Flow)
 """
@@ -38,15 +47,28 @@ logger = logging.getLogger(__name__)
 
 
 def load_config(config_path: str = "config.yaml") -> dict:
+    """Load configuration from YAML file."""
     with open(config_path, "r") as f:
         return yaml.safe_load(f)
 
 
+def get_device(args) -> str:
+    """Resolve device from CLI args or auto-detect."""
+    if args.device:
+        return args.device
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
 def run_2d_experiment(config: dict, args):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    """Run 2D synthetic experiment (NSGF).
+    
+    Reference: Section 5.1, Appendix E.1
+    """
+    device = get_device(args)
     logger.info(f"Running 2D experiment on {device}")
     logger.info(f"Dataset: {config['dataset']}, Steps: {config['sinkhorn']['num_steps']}")
-
+    
+    # Override from args
     if args.dataset:
         config["dataset"] = args.dataset
     if args.steps:
@@ -56,92 +78,143 @@ def run_2d_experiment(config: dict, args):
         config["pool"]["num_batches"] = args.pool_batches
     if args.train_iters:
         config["training"]["num_iterations"] = args.train_iters
-
+    
+    # Setup
     data_loader = DatasetLoader(config)
     model = create_velocity_model_2d(config)
+    
     logger.info(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
-
+    
+    # ---- Training ----
     start_time = time.time()
-    trainer = NSGFTrainer(model=model, data_loader=data_loader, config=config, device=device)
+    
+    trainer = NSGFTrainer(
+        model=model,
+        data_loader=data_loader,
+        config=config,
+        device=device,
+    )
+    
+    # Build trajectory pool
     trainer.build_trajectory_pool()
+    
+    # Train velocity field
     history = trainer.train()
+    
     train_time = time.time() - start_time
     logger.info(f"Training completed in {train_time:.1f}s")
-
+    
+    # ---- Inference ----
     num_eval = config.get("evaluation", {}).get("num_test_samples", 1024)
     num_steps = config.get("inference", {}).get("num_euler_steps", 10)
-    sampler = NSGFSampler(model=model, data_loader=data_loader, num_steps=num_steps, device=device)
+    
+    sampler = NSGFSampler(
+        model=model,
+        data_loader=data_loader,
+        num_steps=num_steps,
+        device=device,
+    )
+    
     samples = sampler.sample(num_eval)
+    
+    # Also get trajectory for visualization
     trajectory = sampler.sample_trajectory(min(200, num_eval))
-
+    
+    # ---- Evaluation ----
     test_samples = data_loader.get_test_samples(num_eval, device)
     evaluator = Evaluation(config, device)
     metrics = evaluator.evaluate(samples, test_samples)
-
+    
     logger.info(f"\n{'='*50}")
     logger.info(f"RESULTS — 2D {config['dataset']}, {num_steps} steps")
     logger.info(f"{'='*50}")
     for k, v in metrics.items():
         logger.info(f"  {k}: {v:.4f}")
     logger.info(f"  Training time: {train_time:.1f}s")
-
+    
+    # ---- Visualization ----
     os.makedirs("results", exist_ok=True)
     plot_2d_samples(
         samples, test_samples,
         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
 
 
 def run_image_experiment(config: dict, args, dataset_name: str):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    """Run image experiment (NSGF++).
+    
+    Reference: Section 5.2, Appendix E.2
+    """
+    device = get_device(args)
     logger.info(f"Running {dataset_name.upper()} experiment on {device}")
-
+    
+    # Override from args
     if args.pool_batches:
         config["pool"]["num_batches"] = args.pool_batches
     if args.train_iters:
         config["nsgf_training"]["num_iterations"] = args.train_iters
         config["nsf_training"]["num_iterations"] = args.train_iters
-
+    
+    # Setup
     data_loader = DatasetLoader(config)
+    
+    # Create models
     nsgf_model = create_velocity_unet(config)
     nsf_model = create_velocity_unet(config)
     phase_predictor = create_phase_predictor(config)
-
+    
     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()):,}")
-
+    
+    # ---- Training ----
     start_time = time.time()
+    
     pp_trainer = NSGFPlusPlusTrainer(
-        nsgf_model=nsgf_model, nsf_model=nsf_model,
-        phase_predictor=phase_predictor, data_loader=data_loader,
-        config=config, device=device,
+        nsgf_model=nsgf_model,
+        nsf_model=nsf_model,
+        phase_predictor=phase_predictor,
+        data_loader=data_loader,
+        config=config,
+        device=device,
     )
+    
     results = pp_trainer.train_all()
+    
     train_time = time.time() - start_time
     logger.info(f"Training completed in {train_time:.1f}s")
-
+    
+    # ---- Inference ----
     inference_cfg = config.get("inference", {})
     nsgf_steps = inference_cfg.get("nsgf_steps", 5)
     nsf_steps = inference_cfg.get("nsf_steps", 55)
     num_gen = config.get("evaluation", {}).get("num_generated", 10000)
-
+    
     sampler = NSGFPlusPlusSampler(
-        nsgf_model=nsgf_model, nsf_model=nsf_model,
-        phase_predictor=phase_predictor, data_loader=data_loader,
-        nsgf_steps=nsgf_steps, nsf_steps=nsf_steps, device=device,
+        nsgf_model=nsgf_model,
+        nsf_model=nsf_model,
+        phase_predictor=phase_predictor,
+        data_loader=data_loader,
+        nsgf_steps=nsgf_steps,
+        nsf_steps=nsf_steps,
+        device=device,
     )
-
+    
     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):
@@ -149,11 +222,15 @@ def run_image_experiment(config: dict, args, dataset_name: str):
         samples = sampler.sample_simple(n)
         all_samples.append(samples.cpu())
     generated = torch.cat(all_samples, dim=0)
-
-    test_samples = data_loader.get_test_samples(num_gen, device="cpu")
+    
+    # ---- Evaluation ----
+    # Get test set
+    eval_loader = data_loader
+    test_samples = eval_loader.get_test_samples(num_gen, device="cpu")
+    
     evaluator = Evaluation(config, device)
     metrics = evaluator.evaluate(generated, test_samples)
-
+    
     logger.info(f"\n{'='*50}")
     logger.info(f"RESULTS — NSGF++ on {dataset_name.upper()}")
     logger.info(f"{'='*50}")
@@ -161,43 +238,60 @@ def run_image_experiment(config: dict, args, dataset_name: str):
         logger.info(f"  {k}: {v:.4f}")
     logger.info(f"  NFE: {nsgf_steps + nsf_steps}")
     logger.info(f"  Training time: {train_time:.1f}s")
-
+    
+    # ---- Visualization ----
     os.makedirs("results", exist_ok=True)
     plot_image_grid(
         generated[:64],
         title=f"NSGF++ — {dataset_name.upper()}",
         save_path=f"results/nsgf_pp_{dataset_name}_samples.png",
     )
+    
+    # Save 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")
     logger.info("Models saved.")
+    
     return metrics
 
 
 def main():
     parser = argparse.ArgumentParser(description="NSGF/NSGF++ Experiments")
-    parser.add_argument("--experiment", type=str, default="2d",
-                        choices=["2d", "mnist", "cifar10"])
-    parser.add_argument("--dataset", type=str, default=None)
-    parser.add_argument("--steps", type=int, default=None)
-    parser.add_argument("--pool-batches", type=int, default=None)
-    parser.add_argument("--train-iters", type=int, default=None)
-    parser.add_argument("--config", type=str, default="config.yaml")
-    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--experiment", type=str, default="2d",
+        choices=["2d", "mnist", "cifar10"],
+        help="Experiment type"
+    )
+    parser.add_argument("--dataset", type=str, default=None, help="2D dataset name")
+    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("--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)")
+    
     args = parser.parse_args()
-
+    
+    # Set seed
     torch.manual_seed(args.seed)
     import numpy as np
     np.random.seed(args.seed)
-
+    
+    # Load config
     full_config = load_config(args.config)
+    
     if args.experiment == "2d":
-        run_2d_experiment(full_config["experiment_2d"], args)
+        config = full_config["experiment_2d"]
+        run_2d_experiment(config, args)
     elif args.experiment == "mnist":
-        run_image_experiment(full_config["experiment_mnist"], args, "mnist")
+        config = full_config["experiment_mnist"]
+        run_image_experiment(config, args, "mnist")
     elif args.experiment == "cifar10":
-        run_image_experiment(full_config["experiment_cifar10"], args, "cifar10")
+        config = full_config["experiment_cifar10"]
+        run_image_experiment(config, args, "cifar10")
     else:
         logger.error(f"Unknown experiment: {args.experiment}")
         sys.exit(1)