| # LEARNING.md — Lessons from NSGF++ Reproduction |
|
|
| All concrete mistakes, war stories, code examples, and derived principles from reproducing arXiv:2401.14069 on Kaggle T4 GPUs. Rules and procedures live in [SKILL.md](SKILL.md). |
|
|
| --- |
|
|
| ## Mistake Catalog |
|
|
| ### #1 — geomloss tensor shape bug (CRITICAL) |
|
|
| **What**: `SamplesLoss` in geomloss requires `(N, D)` or `(B, N, D)` tensors. Image experiments passed `(N, C, H, W)`. |
|
|
| **Impact**: MNIST and CIFAR-10 crash immediately with `ValueError: Input samples 'x' and 'y' should be encoded as (N,D) or (B,N,D) (batch) tensors.` The 2D experiment works fine with shape `(256, 2)`, hiding the bug completely. |
|
|
| **Root cause**: Only tested the 2D code path. Never ran a 10-line test with image-shaped tensors. |
|
|
| **Prevention**: Before building any training loop, test the library with EXACT tensor shapes for every experiment: |
|
|
| ```python |
| from geomloss import SamplesLoss |
| loss_fn = SamplesLoss(loss="sinkhorn", p=2, blur=0.5, potentials=True) |
| |
| # 2D — OK |
| F, G = loss_fn(torch.randn(256, 2, requires_grad=True), torch.randn(256, 2)) |
| |
| # MNIST — CRASHES unless flattened |
| x = torch.randn(128, 1, 28, 28) |
| x_flat = x.view(128, -1).requires_grad_(True) # (128, 784) |
| F, G = loss_fn(x_flat, torch.randn(128, 784)) # ✅ |
| ``` |
|
|
| **Fix applied**: `compute_velocity()` now flattens `(N,C,H,W)→(N,D)` before geomloss, reshapes gradients back after. This pattern applies to all OT libraries (geomloss, POT, ott-jax). |
|
|
| --- |
|
|
| ### #2 — TrajectoryPool re-concatenation every step (MODERATE) |
|
|
| **What**: `torch.cat(self.x_pool, dim=0)` called on the entire pool (512K entries) every single training step. |
|
|
| **Impact**: ~0.5s per step wasted on concatenation vs ~0.05s for the actual forward/backward. Training 10× slower than necessary. |
|
|
| **Root cause**: Didn't profile. The pool was built from a list of tensors and never consolidated. |
|
|
| **Prevention**: Pre-concatenate once after pool building: |
|
|
| ```python |
| def finalize(self): |
| self._all_x = torch.cat(self.x_pool, dim=0) |
| self.x_pool = None # free list memory |
| |
| def sample(self, batch_size, device): |
| idx = torch.randint(0, self._all_x.shape[0], (batch_size,)) |
| return self._all_x[idx].to(device) # O(1) |
| ``` |
|
|
| --- |
|
|
| ### #3 — Incomplete experiment testing (CRITICAL) |
|
|
| **What**: Tested 2D experiments thoroughly on CPU. Shipped MNIST/CIFAR completely untested. |
|
|
| **Impact**: User's first GPU run of MNIST crashes immediately (geomloss bug). User's second run of CIFAR crashes immediately (also geomloss bug + OOM). Two Kaggle sessions wasted. |
|
|
| **Root cause**: 2D success gave false confidence. Different experiment types exercise different code paths — 2D uses `(N, 2)` tensors, images use `(N, C, H, W)`. |
|
|
| **Prevention**: Test EVERY experiment type with `--pool-batches 2 --train-iters 5` before declaring code ready. This takes <60 seconds on CPU. |
|
|
| --- |
|
|
| ### #4 — No checkpointing (MODERATE → CRITICAL at scale) |
|
|
| **What**: No intermediate checkpoints during training. No phase-level saves. |
|
|
| **Impact**: MNIST full run is ~7 hours on T4. Kaggle gives 9 hours per session. Any interruption (timeout, accidental Ctrl+C, OOM partway through) loses everything. CIFAR-10 is impossible in one session without resume. |
|
|
| **Root cause**: Built for "run once and done" — didn't anticipate multi-session training. |
|
|
| **Prevention**: Checkpoint after every phase + every N steps within phases. Implement `--resume-phase`. Test that resume actually loads and skips correctly. |
|
|
| --- |
|
|
| ### #5 — UNet forward pass fragility (LOW-MODERATE) |
|
|
| **What**: `_get_num_res_blocks()` infers block count by dividing module list length by number of levels. Fragile if architecture varies. |
|
|
| **Prevention**: Store `self.num_res_blocks = num_res_blocks` at init. Use it directly in `forward()`. |
|
|
| --- |
|
|
| ### #6 — DataLoader batch size mismatch across phases (CRITICAL) |
|
|
| **What**: `DatasetLoader.sample_target()` lazily creates a DataLoader on first call, caching it. Phase 1 calls with `n=256` → DataLoader has `batch_size=256, drop_last=True`. Phase 2 calls with `n=128` → cached DataLoader still yields 256 → crash: |
|
|
| ``` |
| RuntimeError: The size of tensor a (128) must match the size of tensor b (256) at non-singleton dimension 0 |
| ``` |
|
|
| **Impact**: Phase 2 (NSF) crashes immediately even after Phase 1 completes successfully. The error message doesn't reveal the caching problem — it looks like a model shape issue. |
|
|
| **Root cause**: Lazy initialization without invalidation. Classic stale cache bug. |
|
|
| **Fix applied**: Track `self._image_batch_size` and recreate DataLoader when it changes: |
|
|
| ```python |
| def sample_target(self, n, device="cpu"): |
| if not hasattr(self, "_loader") or self._batch_size != n: |
| self._batch_size = n |
| self._loader = get_image_dataloader(self.dataset_name, batch_size=n) |
| self._iter = iter(self._loader) |
| ``` |
|
|
| --- |
|
|
| ### #7 — CLI flag not overriding all phases (LOW) |
|
|
| **What**: `--train-iters` overrode `nsgf_training.num_iterations` and `nsf_training.num_iterations` but missed `time_predictor.num_iterations` (default 40,000). |
|
|
| **Impact**: Smoke test with `--train-iters 5` completes Phase 1 and 2 in seconds, then hangs for hours on Phase 3. |
|
|
| **Prevention**: Grep config for ALL instances of the parameter: `grep -n num_iterations config.yaml`. Found 3 — must override all 3. |
|
|
| --- |
|
|
| ### #8 — CIFAR-10 Sinkhorn OOM on T4 (CRITICAL) |
|
|
| **What**: Paper's `sinkhorn.batch_size=128` for CIFAR-10. Sinkhorn `tensorized` backend computes 128×128 cost matrix with 3072-dim vectors (flattened 3×32×32). With `potentials=True` + `autograd.grad`, plus 2 calls per flow step × 5 steps per batch = 10 Sinkhorn calls per pool batch. Total: 8+ GB just for Sinkhorn, plus the 38M-param UNet → OOM on T4 16GB. |
|
|
| **Impact**: CIFAR-10 crashes during pool building, before training even starts. |
|
|
| **Root cause**: Used paper hyperparameters verbatim without estimating VRAM for target hardware. Paper authors likely used A100 80GB. |
|
|
| **VRAM math that should have been done upfront**: |
| | Component | Approx VRAM | |
| |-----------|------------| |
| | Sinkhorn call (128×3072, tensorized) | ~600 MB | |
| | × 10 calls per pool batch (with autograd) | ~6+ GB peak | |
| | UNet (38M params, fp32) | ~150 MB | |
| | Gradients + optimizer states | ~450 MB | |
| | **Total** | **~7+ GB** (doesn't account for fragmentation) | |
|
|
| **Fix applied**: Reduce `sinkhorn.batch_size` 128→32, increase `pool.num_batches` 2500→10000. Total pool entries unchanged (1.6M). Added `--sinkhorn-batch` CLI flag. |
|
|
| --- |
|
|
| ### #9 — No GPU memory freed between phases (MODERATE) |
|
|
| **What**: After pool building finishes, Sinkhorn computation graph CUDA allocations stay cached. Training starts with reduced available VRAM. |
|
|
| **Fix applied**: `torch.cuda.empty_cache()` after pool building + `del pool` after finalization. |
|
|
| --- |
|
|
| ### #10 — Multi-GPU assumption (LOW) |
|
|
| **What**: User has T4×2 on Kaggle. Code is single-GPU. Second GPU sits idle. |
|
|
| **Prevention**: Document explicitly: "Single-GPU only. T4×2 wastes the second GPU. Use single T4 instead, or add DDP." |
|
|
| --- |
|
|
| ## Principles |
|
|
| Derived from the mistakes above. Each traces back to one or more specific failures. |
|
|
| 1. **Read the appendix first.** The appendix is the recipe; the main paper is the story. *(from: all mistakes — hyperparams were in appendix)* |
|
|
| 2. **Test the boundaries, not just the happy path.** The bug is always in the path you didn't test. *(from: #1, #3)* |
|
|
| 3. **Library APIs are opaque until tested.** Don't assume a function accepts your tensor shape just because it "makes sense." Write a 10-line test script. *(from: #1)* |
|
|
| 4. **Pre-concatenate, don't re-concatenate.** Any data structure built once and sampled many times should be finalized into a single tensor. *(from: #2)* |
|
|
| 5. **The user's time is more expensive than your time.** A crash on their GPU after 5 minutes of setup is worse than you spending 30 extra minutes testing. *(from: #1, #3, #8)* |
|
|
| 6. **Flatten for OT libraries.** geomloss, POT, ott-jax all expect `(N, D)` point clouds. Images must be flattened. #1 gotcha in OT-based generative models. *(from: #1)* |
|
|
| 7. **Store training state on CPU, compute on GPU.** Pools and replay buffers on CPU; only minibatch goes to GPU. *(from: #8, #9)* |
|
|
| 8. **Multi-phase training = multiple separate trainers.** Each phase gets its own optimizer. Previous phase's model → `eval()`. *(from: #6)* |
|
|
| 9. **Shared objects across phases are landmines.** Cached DataLoaders, iterators, batch sizes — any phase-specific param can silently break later phases. *(from: #6)* |
|
|
| 10. **CLI overrides must be exhaustive.** N copies of a parameter in config → CLI must touch all N. *(from: #7)* |
|
|
| 11. **Paper hyperparameters assume paper hardware.** Re-derive batch sizes from VRAM constraints. Keep total samples seen constant. *(from: #8)* |
|
|
| 12. **Estimate VRAM before running, not after OOM.** Sinkhorn: O(N² × D). Model: params × 4 bytes × 3. Write it down before the first GPU run. *(from: #8)* |
|
|
| 13. **Checkpoint at phase boundaries, not just step boundaries.** Phase-level checkpoints enable `--resume-phase`. Step-level checkpoints within long phases are a bonus. *(from: #4)* |
|
|
| 14. **Free GPU memory between phases.** `torch.cuda.empty_cache()` + `del` large objects between phases with different memory patterns. *(from: #9)* |
|
|
| 15. **Document what your code does NOT support.** Single-GPU? No mixed precision? Say so explicitly. *(from: #10)* |
|
|
| --- |
|
|
| ## Sinkhorn VRAM Reference Table |
|
|
| For quick VRAM estimation during paper reproduction planning. `tensorized` backend, fp32, single `SamplesLoss` call with `potentials=True`. |
|
|
| | N (batch) | D (dim) | Example | Approx VRAM/call | |
| |-----------|---------|---------|-----------------| |
| | 256 | 2 | 2D points | ~1 MB | |
| | 256 | 784 | MNIST 28×28 | ~200 MB | |
| | 128 | 3072 | CIFAR 3×32×32 | ~600 MB | |
| | 32 | 3072 | CIFAR (T4-safe) | ~40 MB | |
|
|
| Pool building does ~10 calls per batch (2 potentials × 5 flow steps). Multiply accordingly. |
|
|
| --- |
|
|
| ## Timeline of Bugs |
|
|
| | When | What broke | How discovered | Sessions wasted | |
| |------|-----------|----------------|-----------------| |
| | Initial ship | geomloss shape (MNIST/CIFAR) | User's first GPU run | 1 | |
| | After geomloss fix | DataLoader batch mismatch (Phase 2) | My CPU test | 0 (caught in sandbox) | |
| | After DataLoader fix | `--train-iters` missed Phase 3 | My CPU test | 0 (caught in sandbox) | |
| | User's full MNIST run | No checkpointing, interrupted | User's Kaggle timeout | 1 | |
| | User's CIFAR run | Sinkhorn OOM on T4 | User's Kaggle log | 1 | |
| | After OOM fix | No `empty_cache()` between phases | Analysis of VRAM patterns | 0 (preventive) | |
|
|
| **Total Kaggle sessions wasted by user: 3.** The bugs caught in sandbox (0 sessions wasted) vs bugs caught on user hardware (3 sessions) reinforce principle #5: test everything before shipping. |
|
|
