diff --git a/README.md b/README.md index 7898d54..82826e6 100644 --- a/README.md +++ b/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 | diff --git a/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin index b67a2e9..7a68cec 100644 --- a/conf/gtauav_balanced.gin +++ b/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 diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index f0ecd2e..2b2754d 100644 --- a/src/training/train_gtauav.py +++ b/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. - steps_per_epoch = len(train_loader) + # Scheduler — cosine with linear warmup (counted in optimizer steps). + 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: - scaler.unscale_(optimizer) - nn.utils.clip_grad_norm_( - model.trainable_parameters(), - max_norm=cfg.grad_clip, - ) + # Optimizer step only after accumulating `accum` micro-batches. + is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(train_loader) + if is_accum_step: + if cfg.grad_clip > 0: + scaler.unscale_(optimizer) + nn.utils.clip_grad_norm_( + model.trainable_parameters(), + max_norm=cfg.grad_clip, + ) - # --- Gradient monitoring (after unscale, before step) --- - if cfg.log_grad_norms and n_batches % 50 == 0: - grad_norms = compute_gradient_norms(model, loss_fn) - tracker.log_gradients(epoch, grad_norms, step=global_step) - if n_batches == 0: - log_gradient_summary(grad_norms) + # --- Gradient monitoring (after unscale, before step) --- + if cfg.log_grad_norms and n_batches % (50 * accum) < accum: + grad_norms = compute_gradient_norms(model, loss_fn) + tracker.log_gradients(epoch, grad_norms, step=global_step) + if n_batches < accum: + log_gradient_summary(grad_norms) - scaler.step(optimizer) - scaler.update() - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*") - scheduler.step() + scaler.step(optimizer) + scaler.update() + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.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: