Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add gradient accumulation support

Browse Source 
- 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>
This commit is contained in:
pikaliov
2026-04-21 21:21:56 +03:00
parent b5c8015616
commit 46b1208891
3 changed files with 70 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
README.md
View File

@@ -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 |

1
conf/gtauav_balanced.gin
View File

@@ -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

71
src/training/train_gtauav.py
View File

@@ -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: