From 772acd8f0d285e0dc323089c1f92457a4ff9bcba Mon Sep 17 00:00:00 2001 From: Yann Ahlgrim Date: Thu, 18 Jun 2026 16:25:52 +0200 Subject: [PATCH] gradient accumulation --- configs/grids/lp_grid.yaml | 3 ++- src/train_supervised.py | 10 +++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/configs/grids/lp_grid.yaml b/configs/grids/lp_grid.yaml index b3732ed..bd6cf12 100644 --- a/configs/grids/lp_grid.yaml +++ b/configs/grids/lp_grid.yaml @@ -9,7 +9,8 @@ grid: optimization.weight_decay: [5.0e-4, 0] optimization.momentum: [0.9] optimization.lr_schedule: [cosine] - data.batch_size: [512, 1024] + optimization.gradient_accumulation_steps: [8, 16] + data.batch_size: [64] meta.representation_type: [last_avgpool, last4_avgpool_concat] meta.head_type: [linear, bn_linear] diff --git a/src/train_supervised.py b/src/train_supervised.py index eb3eb20..624fb90 100644 --- a/src/train_supervised.py +++ b/src/train_supervised.py @@ -213,6 +213,7 @@ def main(args, resume_preempt=False): l_args = args["logging"] v_args = args["validation"] es_args = o_args["early_stopping"] + accum_steps = o_args.get("gradient_accumulation_steps", 1) folder = resolve_log_dir(args, stage="train") tag = l_args["write_tag"] @@ -431,6 +432,7 @@ def main(args, resume_preempt=False): loss_meter = AverageMeter() + optimizer.zero_grad() for itr, (imgs, labels) in enumerate(train_loader): imgs = imgs.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) @@ -441,9 +443,11 @@ def main(args, resume_preempt=False): outputs = model(imgs) loss = criterion(outputs, labels) - optimizer.zero_grad() - loss.backward() - optimizer.step() + (loss / accum_steps).backward() + + if (itr + 1) % accum_steps == 0: + optimizer.step() + optimizer.zero_grad() loss_meter.update(loss.item(), n=labels.size(0))