Upload folder using huggingface_hub

train_unet.py

@@ -83,6 +83,13 @@ def parse_args() -> argparse.Namespace:
         help="YAML with UNet + CFM hyperparameters (default: unet_config.yaml next to this script)",
     )
 
+    p.add_argument(
+        "--data-percent",
+        type=int,
+        default=100,
+        choices=[10, 20, 30, 60, 80, 100],
+        help="Use only this percentage of the (possibly filtered) training dataset.",
+    )
     return p.parse_args()
 
 
@@ -257,13 +264,28 @@ def main() -> None:
         [
             v2.ToTensor(),
             v2.ToDtype(torch.float32, scale=True),
-            v2.Resize((64,64)),
+            v2.Resize((32,32)),
             v2.Normalize(mean=[0.0, 0.0, 0.0], std=[1.0, 1.0, 1.0]),
         ]
     )
     train_dataset = load_training_dataset(args, transforms)
     print(f"Dataset: {args.dataset}, size={len(train_dataset)}")
 
+    orig_len = len(train_dataset)
+    if args.data_percent < 100:
+        new_len = max(1, int(orig_len * args.data_percent / 100.0))
+
+        g = torch.Generator()
+        g.manual_seed(args.seed)
+
+        perm = torch.randperm(orig_len, generator=g)
+        indices = perm[:new_len].tolist()
+        torch.save(perm[:new_len], os.path.join(args.save_dir, "indices.pt"))
+        train_dataset = Subset(train_dataset, indices)
+        print(f"Subsampled dataset: {args.data_percent}% -> {len(train_dataset)} samples")
+    else:
+        print(f"Using full dataset: {orig_len} samples")
+        
     dummy_dataloader = DataLoader(
         train_dataset,
         batch_size=args.batch_size,

unet_cifar.yaml

@@ -11,7 +11,7 @@ weight_decay: 0.0
 # NeuralODE visualization / sampling
 save_ep: 30
 inference_steps: 100
-vis_batch_size: 4
+vis_batch_size: 8
 
 # UNet (torchcfm UNetModelWrapper)
 num_res_blocks: 2