def train()

in PyTorchClassification/train.py [0:0]

• 72 lines of code
• 12 McCabe index (conditional complexity)

def train(train_loader, model, criterion, optimizer, epoch, param_copy = None):

    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top3 = AverageMeter()
    top5 = AverageMeter()

    # switch to train mode
    model.train()

    end = time.time()
    print_log('Epoch:{0}'.format(epoch))
    print_log('Itr\t\tTime\t\tData\t\tLoss\t\tPrec@1\t\tPrec@3\t\tPrec@5')
    for param_group in optimizer.param_groups:
        writer.add_scalar('training/learning_rate', param_group['lr'] / param_group['lr_mult'],
                             len(train_loader)*epoch)
        break

    # Make lr really low for the first couple iterations
    iterations_processed = 0
    if args.warm_up_iterations:
        print_log('Warming up the training for {} iterations with a very small learning rate'.format(args.warm_up_iterations))
        adjust_learning_rate(optimizer, 10000)

    data_iterator = tqdm.tqdm(enumerate(train_loader), total=len(train_loader))
    for i, (input, im_id, target) in data_iterator:
        iterations_processed = iterations_processed + 1
        if args.warm_up_iterations and iterations_processed > args.warm_up_iterations:
            args.warm_up_iterations = 0
            print_log('Finished warm-up phase')
            adjust_learning_rate(optimizer, epoch)

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

        input = input.cuda()

        target = target.cuda()
        input_var = torch.autograd.Variable(input)
        target_var = torch.autograd.Variable(target)

        # compute output
        output = model(input_var)
        loss = criterion(output, target_var)

        # measure accuracy and record loss
        prec1, prec3, prec5 = accuracy(output.data, target, topk=(1, 3, 5))

        if args.distributed:
            reduced_loss = reduce_tensor(loss.data)
            prec1 = reduce_tensor(prec1)
            prec5 = reduce_tensor(prec5)
        else:
            reduced_loss = loss.data

        losses.update(reduced_loss.item(), input.size(0))
        top1.update(prec1.item(), input.size(0))
        top3.update(prec3.item(), input.size(0))
        top5.update(prec5.item(), input.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        if args.fp16:
            optimizer.backward(loss)
        else:
            loss.backward()
        optimizer.step()
        torch.cuda.synchronize()

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

        writer.add_scalar('training/loss', losses.val, len(train_loader)*epoch + i)
        writer.add_scalars('training/topk', {'top1':top1.val,
                                                   'top3':top3.val,
                                                   'top5':top5.val}, 
                                                 len(train_loader)*epoch + i)
        if i % args.print_freq == 0:
            x = vutils.make_grid(input[0]/2 + 0.5)
            writer.add_image('Preprocessed training images', x, len(train_loader)*epoch + i)
            #for name, param in model.named_parameters():
            #    writer.add_histogram(name, param.clone().cpu().data.numpy(), i)

            print_log('[{0}/{1}]\t'
                '{batch_time.val:.2f} ({batch_time.avg:.2f})\t'
                '{data_time.val:.2f} ({data_time.avg:.2f})\t'
                '{loss.val:.3f} ({loss.avg:.3f})\t'
                '{top1.val:.2f} ({top1.avg:.2f})\t'
                '{top3.val:.2f} ({top3.avg:.2f})\t'
                '{top5.val:.2f} ({top5.avg:.2f})'.format(i, len(train_loader), batch_time=batch_time,
                data_time=data_time, loss=losses, top1=top1, top3=top3, top5=top5))

    args.warm_up_iterations = 0
    print_log(' *** Training summary at epoch {epoch:d}: Prec@1 {top1.avg:.3f} '.format(epoch=epoch, top1=top1) +
              'Prec@3 {top3.avg:.3f} Prec@5 {top5.avg:.3f} Loss {loss.avg:.3f}'.format(top3=top3, top5=top5, loss=losses))