Add gradient accumulation support

- New config field grad_accum_steps (default=1, no change in behavior)
- Loss scaled by 1/accum, optimizer step every N micro-batches
- Scheduler counts optimizer steps (not micro-batches)
- CLI flag --grad-accum for override
- Document gradient accumulation and in-batch negatives in README

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

README.md

@@ -186,6 +186,30 @@ Symmetric InfoNCE with learnable temperature (CLIP-style `logit_scale`):
 
 Loss and adapters run in **fp32** (AMP autocast disabled) to prevent gradient overflow.
 
+### Gradient accumulation
+
+With `batch_size=8` on a 24 GB GPU, VRAM is the bottleneck. Gradient accumulation
+emulates a larger effective batch without extra memory:
+
+```
+effective_batch_size = batch_size × grad_accum_steps
+```
+
+| Setting | `batch_size` | `grad_accum_steps` | Effective batch | In-batch negatives |
+|---------|:---:|:---:|:---:|:---:|
+| Default | 8 | 1 | 8 | 7 |
+| Recommended | 8 | 8 | 64 | 7 per micro-batch |
+
+**Note:** gradient accumulation averages gradients across micro-batches, but each
+micro-batch still only sees `batch_size` in-batch negatives. To increase the number
+of negatives per forward pass, increase `batch_size` directly (requires more VRAM).
+
+```bash
+# Example: effective batch of 64 with 8 accumulation steps
+python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \
+    --filter-meta meta/seg_filter.json --grad-accum 8
+```
+
 ### Metrics
 
 | Metric | Formula | Direction |

conf/gtauav_balanced.gin

@@ -15,6 +15,7 @@ TrainConfigGTAUAV.learning_rate = 1e-4
 TrainConfigGTAUAV.text_lr_factor = 0.1
 TrainConfigGTAUAV.weight_decay = 1e-4
 TrainConfigGTAUAV.grad_clip = 1.0
+TrainConfigGTAUAV.grad_accum_steps = 1
 TrainConfigGTAUAV.use_amp = True
 TrainConfigGTAUAV.eval_every = 2
 TrainConfigGTAUAV.warmup_epochs = 2

src/training/train_gtauav.py

@@ -84,6 +84,7 @@ class TrainConfigGTAUAV:
     text_lr_factor: float = 0.1  # text encoder LR = learning_rate * factor
     weight_decay: float = 1e-4
     grad_clip: float = 1.0
+    grad_accum_steps: int = 1  # gradient accumulation steps (effective_batch = batch_size * accum)
     use_amp: bool = True
     eval_every: int = 2
     warmup_epochs: int = 2
@@ -418,7 +419,11 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         pin_memory=True,
     )
 
-    LOGGER.info("train=%d  test=%d  batch=%d", len(train_ds), len(test_ds), cfg.batch_size)
+    effective_batch = cfg.batch_size * cfg.grad_accum_steps
+    LOGGER.info(
+        "train=%d  test=%d  batch=%d  accum=%d  effective_batch=%d",
+        len(train_ds), len(test_ds), cfg.batch_size, cfg.grad_accum_steps, effective_batch,
+    )
 
     # Optimizer — per-group LR (text encoder gets lower LR).
     param_groups = _build_param_groups(model, cfg.learning_rate, cfg.text_lr_factor)
@@ -433,8 +438,8 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         lr_info += f"  text={cfg.learning_rate * cfg.text_lr_factor:.0e}"
     LOGGER.info("Optimizer: AdamW  LR: %s  warmup=%d epochs", lr_info, cfg.warmup_epochs)
 
-    # Scheduler — cosine with linear warmup.
+    # Scheduler — cosine with linear warmup (counted in optimizer steps).
-    steps_per_epoch = len(train_loader)
+    steps_per_epoch = math.ceil(len(train_loader) / cfg.grad_accum_steps)
     total_steps = cfg.epochs * steps_per_epoch
     warmup_steps = cfg.warmup_epochs * steps_per_epoch
     with warnings.catch_warnings():
@@ -488,8 +493,11 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             unit="batch",
             leave=False,
         )
+        accum = cfg.grad_accum_steps
         for batch in pbar:
-            optimizer.zero_grad(set_to_none=True)
+            # Zero gradients only at the start of each accumulation window.
+            if n_batches % accum == 0:
+                optimizer.zero_grad(set_to_none=True)
 
             drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
             sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
@@ -516,32 +524,38 @@ def train(cfg: TrainConfigGTAUAV) -> None:
                 total_epochs=cfg.epochs,
             )
 
-            total_loss = loss_dict["total"]
+            # Scale loss by accumulation steps so gradients average correctly.
+            total_loss = loss_dict["total"] / accum
             scaler.scale(total_loss).backward()
 
-            if cfg.grad_clip > 0:
+            # Optimizer step only after accumulating `accum` micro-batches.
-                scaler.unscale_(optimizer)
+            is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(train_loader)
-                nn.utils.clip_grad_norm_(
+            if is_accum_step:
-                    model.trainable_parameters(),
+                if cfg.grad_clip > 0:
-                    max_norm=cfg.grad_clip,
+                    scaler.unscale_(optimizer)
-                )
+                    nn.utils.clip_grad_norm_(
+                        model.trainable_parameters(),
+                        max_norm=cfg.grad_clip,
+                    )
 
-            # --- Gradient monitoring (after unscale, before step) ---
+                # --- Gradient monitoring (after unscale, before step) ---
-            if cfg.log_grad_norms and n_batches % 50 == 0:
+                if cfg.log_grad_norms and n_batches % (50 * accum) < accum:
-                grad_norms = compute_gradient_norms(model, loss_fn)
+                    grad_norms = compute_gradient_norms(model, loss_fn)
-                tracker.log_gradients(epoch, grad_norms, step=global_step)
+                    tracker.log_gradients(epoch, grad_norms, step=global_step)
-                if n_batches == 0:
+                    if n_batches < accum:
-                    log_gradient_summary(grad_norms)
+                        log_gradient_summary(grad_norms)
 
-            scaler.step(optimizer)
+                scaler.step(optimizer)
-            scaler.update()
+                scaler.update()
-            with warnings.catch_warnings():
+                with warnings.catch_warnings():
-                warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
+                    warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
-                scheduler.step()
+                    scheduler.step()
+                global_step += 1
 
-            # --- Per-step tracking ---
+            # --- Per-batch tracking (log unscaled loss) ---
+            raw_loss = total_loss.item() * accum  # undo /accum for logging
             step_metrics = {
-                "loss": float(total_loss.item()),
+                "loss": raw_loss,
                 "temperature": float(loss_dict["temperature"].item()),
                 "gate_q": float(loss_dict["gate_q"].item()),
                 "gate_g": float(loss_dict["gate_g"].item()),
@@ -553,10 +567,9 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             for key, val in loss_dict.items():
                 agg[key] = agg.get(key, 0.0) + float(val.item())
             n_batches += 1
-            global_step += 1
 
             pbar.set_postfix(
-                loss=f"{total_loss.item():.3f}",
+                loss=f"{raw_loss:.3f}",
                 tau=f"{loss_dict['temperature'].item():.4f}",
                 gq=f"{loss_dict['gate_q'].item():.3f}",
                 gg=f"{loss_dict['gate_g'].item():.3f}",
@@ -726,6 +739,10 @@ def main() -> None:
         "--batch-size", type=int, default=None,
         help="Batch size.",
     )
+    parser.add_argument(
+        "--grad-accum", type=int, default=None,
+        help="Gradient accumulation steps (effective_batch = batch_size * accum).",
+    )
     parser.add_argument(
         "--epochs", type=int, default=None,
         help="Number of epochs.",
@@ -775,6 +792,8 @@ def main() -> None:
         cfg.resume_from = args.resume
     if args.batch_size is not None:
         cfg.batch_size = args.batch_size
+    if args.grad_accum is not None:
+        cfg.grad_accum_steps = args.grad_accum
     if args.epochs is not None:
         cfg.epochs = args.epochs
     if args.lr is not None: