| 205 | |
| 206 | |
| 207 | def train_epoch(*, model, optimizer, data_loader, scheduler, criterion, mixup_fn, |
| 208 | accelerator: Accelerator, epoch, config): |
| 209 | model.train() |
| 210 | |
| 211 | num_steps = len(data_loader) |
| 212 | batch_time = AverageMeter() |
| 213 | model_time = AverageMeter() |
| 214 | loss_meter = AverageMeter() |
| 215 | |
| 216 | end = time.time() |
| 217 | |
| 218 | gradient_accumulation_steps = config.TRAIN.ACCUMULATION_STEPS |
| 219 | |
| 220 | for step, (samples, targets) in enumerate(data_loader): |
| 221 | iter_begin_time = time.time() |
| 222 | |
| 223 | if mixup_fn is not None: |
| 224 | samples, targets = mixup_fn(samples, targets) |
| 225 | |
| 226 | with accelerator.accumulate(model): |
| 227 | outputs = model(samples) |
| 228 | loss = criterion(outputs, targets) |
| 229 | accelerator.backward(loss) |
| 230 | if accelerator.sync_gradients: |
| 231 | accelerator.clip_grad_norm_(model.parameters(), config.TRAIN.CLIP_GRAD) |
| 232 | optimizer.step() |
| 233 | optimizer.zero_grad() |
| 234 | |
| 235 | accelerator.wait_for_everyone() |
| 236 | |
| 237 | if (step + 1) % gradient_accumulation_steps == 0: |
| 238 | if scheduler is not None: |
| 239 | scheduler.step_update((epoch * num_steps + step) // gradient_accumulation_steps) |
| 240 | |
| 241 | batch_time.update(time.time() - end) |
| 242 | model_time.update(time.time() - iter_begin_time) |
| 243 | loss_meter.update(loss.item()) |
| 244 | end = time.time() |
| 245 | |
| 246 | if accelerator.is_main_process and step % config.PRINT_FREQ == 0: |
| 247 | lr = optimizer.param_groups[0]['lr'] |
| 248 | memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) |
| 249 | etas = batch_time.avg * (num_steps - step) |
| 250 | |
| 251 | logger.info( |
| 252 | f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{step}/{num_steps}]\t' |
| 253 | f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.10f}\t' |
| 254 | f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t' |
| 255 | f'model_time {model_time.val:.4f} ({model_time.avg:.4f})\t' |
| 256 | f'loss {loss_meter.val:.8f} ({loss_meter.avg:.4f})\t' |
| 257 | f'mem {memory_used:.0f}MB') |
| 258 | |
| 259 | |
| 260 | @torch.no_grad() |