def train()

in main_byol.py [0:0]

• 46 lines of code
• 10 McCabe index (conditional complexity)

def train(train_loader, model, optimizer, epoch, args):
    batch_time = utils.AverageMeter('Time', ':6.3f')
    data_time = utils.AverageMeter('Data', ':6.3f')
    losses = utils.AverageMeter('Loss', ':.4e')
    curr_lr = utils.InstantMeter('LR', '')
    curr_mom = utils.InstantMeter('MOM', '')
    progress = utils.ProgressMeter(
        len(train_loader),
        [curr_lr, curr_mom, batch_time, data_time, losses],
        prefix="Epoch: [{}/{}]\t".format(epoch, args.epochs))

    # iter info
    batch_iter = len(train_loader)
    max_iter = float(batch_iter * args.epochs)

    # switch to train mode
    model.train()
    if "EMAN" in args.arch:
        print("setting the key model to eval mode when using EMAN")
        if hasattr(model, 'module'):
            model.module.target_net.eval()
        else:
            model.target_net.eval()

    end = time.time()
    for i, (images, _) in enumerate(train_loader):
        # update model momentum
        curr_iter = float(epoch * batch_iter + i)

        # measure data loading time
        data_time.update(time.time() - end)

        if args.gpu is not None:
            images[0] = images[0].cuda(args.gpu, non_blocking=True)
            images[1] = images[1].cuda(args.gpu, non_blocking=True)

        # warmup learning rate
        if epoch < args.warmup_epoch:
            warmup_step = args.warmup_epoch * batch_iter
            curr_step = epoch * batch_iter + i + 1
            lr_schedule.warmup_learning_rate(optimizer, curr_step, warmup_step, args)
        curr_lr.update(optimizer.param_groups[0]['lr'])

        # compute loss
        loss = model(im_v1=images[0], im_v2=images[1])

        # measure accuracy and record loss
        losses.update(loss.item(), images[0].size(0))

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

        # update the target model
        if hasattr(model, 'module'):
            model.module.momentum_update(curr_iter, max_iter)
            curr_mom.update(model.module.curr_m)
        else:
            model.momentum_update(curr_iter, max_iter)
            curr_mom.update(model.curr_m)

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:
            progress.display(i)