Function train_one_epoch_distill_using_saved_logits

TinyViT/main.py:284–400

(args, config, model, criterion, data_loader, optimizer, epoch, mixup_fn, lr_scheduler, loss_scaler)

def train_one_epoch_distill_using_saved_logits(args, config, model, criterion, data_loader, optimizer, epoch, mixup_fn, lr_scheduler, loss_scaler):
    model.train()
    set_bn_state(config, model)
    optimizer.zero_grad()

    num_steps = len(data_loader)
    batch_time = AverageMeter()
    loss_meter = AverageMeter()
    norm_meter = AverageMeter()
    scaler_meter = AverageMeter()
    meters = defaultdict(AverageMeter)

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

    num_classes = config.MODEL.NUM_CLASSES
    topk = config.DISTILL.LOGITS_TOPK

    for idx, ((samples, targets), (logits_index, logits_value, seeds)) in enumerate(data_loader):
        normal_global_idx = epoch * NORM_ITER_LEN + \
            (idx * NORM_ITER_LEN // num_steps)

        samples = samples.cuda(non_blocking=True)
        targets = targets.cuda(non_blocking=True)

        if mixup_fn is not None:
            samples, targets = mixup_fn(samples, targets, seeds)
            original_targets = targets.argmax(dim=1)
        else:
            original_targets = targets
        meters['data_time'].update(time.time() - data_tic)

        with torch.cuda.amp.autocast(enabled=config.AMP_ENABLE):
            outputs = model(samples)

        # recover teacher logits
        logits_index = logits_index.long()
        logits_value = logits_value.float()
        logits_index = logits_index.cuda(non_blocking=True)
        logits_value = logits_value.cuda(non_blocking=True)
        minor_value = (1.0 - logits_value.sum(-1, keepdim=True)
                       ) / (num_classes - topk)
        minor_value = minor_value.repeat_interleave(num_classes, dim=-1)
        outputs_teacher = minor_value.scatter_(-1, logits_index, logits_value)

        loss = criterion(outputs, outputs_teacher)
        loss = loss / config.TRAIN.ACCUMULATION_STEPS

        # this attribute is added by timm on one optimizer (adahessian)
        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)
        if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0:
            optimizer.zero_grad()
            lr_scheduler.step_update(