Function train

main.py:194–238

(loader, model, crit, opt, epoch, labels, label_update=True)

        lr_scheduler.step()

def train(loader, model, crit, opt, epoch, labels, label_update=True):
    batch_time = AverageMeter()
    losses = AverageMeter()
    data_time = AverageMeter()
    forward_time = AverageMeter()
    backward_time = AverageMeter()

    # switch to train mode
    model.train()

    end = time.time()
    for i, (inputs, targets, index) in enumerate(loader):
        data_time.update(time.time() - end)
        inputs, targets = inputs.cuda(), targets.cuda()
        outputs = model(inputs)

        # labels updated
        if label_update:
            predict = outputs.argmax(dim=1).data.cpu().numpy()
            for index_item, pred in zip(index, predict):
                labels[index_item] = int(pred)

        loss = crit(outputs, targets)

        # record loss
        losses.update(loss.item(), inputs.size(0))

        # compute gradient and do SGD step
        opt.zero_grad()
        loss.backward()
        opt.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        if args.verbose and (i % 100) == 0:
            logger.info('Epoch: [{0}][{1}/{2}]\t'
                  'LR: {3}\t'
                  'Time: {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data: {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss: {loss.val:.4f} ({loss.avg:.4f})'
                  .format(epoch, i, len(loader), opt.param_groups[0]['lr'], batch_time=batch_time,
                          data_time=data_time, loss=losses))

    return losses.avg

def compute_labels(dataloader, model, class_num, data_len, init_via_forward=False):
    '''Pre-generate pseudo labels via network forward or uniformly assignment'''