Function train_one_epoch

classification/main.py:403–535

(config,
                    model,
                    criterion,
                    data_loader,
                    optimizer,
                    epoch,
                    mixup_fn,
                    lr_scheduler,
                    amp_autocast=suppress,
                    loss_scaler=None,
                    model_ema=None)

def train_one_epoch(config,
                    model,
                    criterion,
                    data_loader,
                    optimizer,
                    epoch,
                    mixup_fn,
                    lr_scheduler,
                    amp_autocast=suppress,
                    loss_scaler=None,
                    model_ema=None):
    model.train()
    optimizer.zero_grad()

    num_steps = len(data_loader)
    batch_time = AverageMeter()
    model_time = AverageMeter()
    loss_meter = AverageMeter()
    norm_meter = MyAverageMeter(300)

    start = time.time()
    end = time.time()

    amp_type = torch.float16 if config.AMP_TYPE == 'float16' else torch.bfloat16
    for idx, (samples, targets) in enumerate(data_loader):
        iter_begin_time = time.time()
        samples = samples.cuda(non_blocking=True)
        targets = targets.cuda(non_blocking=True)

        if mixup_fn is not None:
            samples, targets = mixup_fn(samples, targets)

        if not obsolete_torch_version(TORCH_VERSION,
                                      (1, 9)) and config.AMP_OPT_LEVEL != 'O0':
            with amp_autocast(dtype=amp_type):
                outputs = model(samples)
        else:
            with amp_autocast():
                outputs = model(samples)

        if config.TRAIN.ACCUMULATION_STEPS > 1:
            if not obsolete_torch_version(
                    TORCH_VERSION, (1, 9)) and config.AMP_OPT_LEVEL != 'O0':
                with amp_autocast(dtype=amp_type):
                    loss = criterion(outputs, targets)
                    loss = loss / config.TRAIN.ACCUMULATION_STEPS
            else:
                with amp_autocast():
                    loss = criterion(outputs, targets)
                    loss = loss / config.TRAIN.ACCUMULATION_STEPS
            if config.AMP_OPT_LEVEL != 'O0':
                is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
                grad_norm = loss_scaler(loss,
                                        optimizer,
                                        clip_grad=config.TRAIN.CLIP_GRAD,
                                        parameters=model.parameters(),
                                        create_graph=is_second_order,
                                        update_grad=(idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0)