def train()

in Classification/main_imagenet.py [0:0]

• 47 lines of code
• 17 McCabe index (conditional complexity)

def train(train_loader, model, criterion, optimizer, epoch, args):
    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    losses = AverageMeter('Loss', ':.4e')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    progress = ProgressMeter(len(train_loader), batch_time, data_time, losses, top1,
                             top5, prefix="Epoch: [{}]".format(epoch))

    # switch to train mode
    model.train()

    end = time.time()
    for i, (input, target) in enumerate(train_loader):
        if args.test_run and i>args.test_iter:
            break

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

        if args.gpu is not None:
            input = input.cuda(args.gpu, non_blocking=True)
        target = target.cuda(args.gpu, non_blocking=True)

        # compute output
        output = model(input)
        loss = criterion(output, target)

        # measure accuracy and record loss
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        losses.update(loss.item(), input.size(0))
        top1.update(acc1[0], input.size(0))
        top5.update(acc5[0], input.size(0))

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

        if args.lr_scheduler == 'cosine':
            args.current_scheduler.step()

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

        global n_iter
        n_iter = n_iter + 1

        if i % args.print_freq == 0:
            progress.print(i)
            if args.tensorboard:
                if not args.multiprocessing_distributed or (args.multiprocessing_distributed
                    and args.rank % args.ngpus_per_node == 0):
                    args.writer.add_scalar('CurrentLoss', losses.val, n_iter)



    if args.tensorboard:
        if not args.multiprocessing_distributed or (args.multiprocessing_distributed
                    and args.rank % args.ngpus_per_node == 0):
            args.writer.add_scalar('Loss/Train',losses.avg,epoch+1)
            args.writer.add_scalar('Prec/Train1',top1.avg,epoch+1)
            args.writer.add_scalar('Prec/Train5', top5.avg, epoch + 1)

    args.train_losses.append(losses.avg)
    args.train_top1.append(top1.avg)
    args.train_top5.append(top5.avg)