nbansal90 / facial_attribute_segmentation Goto Github PK

            # training
            training_top1 = AverageMeter()
            training_losses = AverageMeter()
            training_batch_time = AverageMeter()
            model.train()
            end = time.time()
            for i, (name, input_img, target) in enumerate(train_loader):
                    target = target.cuda(async=True)
                    input_img = input_img.float().cuda()
                    input_var = torch.autograd.Variable(input_img)
                    target_var = torch.autograd.Variable(target)
                    output = model(input_var)
                    #output = output.permute(0,2,3,1)
                    #output = torch.max(output,dim=3)[0]
                    loss = criterion(output, target_var)

                    prec1 = accuracy(output.data, target, topk=(1,))[0]
                    training_losses.update(loss.item(), input_img.size(0))
                    training_top1.update(prec1.item(), input_img.size(0))
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()

                    training_batch_time.update(time.time() - end)
                    end = time.time()

                    print('Epoch: [{0}][{1}/{2}]\t'
                            'Time {training_batch_time.val:.3f} ({training_batch_time.avg:.3f})\t'
                            'Loss {training_losses.val:.4f} ({training_losses.avg:.4f})\t'
                            'Prec@1 {training_top1.val:.3f} ({training_top1.avg:.3f})'.format(
                                    epoch, i, len(train_loader), training_batch_time=training_batch_time,
                                    training_losses=training_losses, training_top1=training_top1))
            # evaluation

            eval_top1 = AverageMeter()
            eval_losses = AverageMeter()
            model.eval()
            with torch.no_grad():
                    for i, (name, input_img, target) in enumerate(val_loader):
                            target = target.cuda(async=True)
                            input_img = input_img.float().cuda()
                           input_var = torch.autograd.Variable(input_img)
                            target_var = torch.autograd.Variable(target)
                            output = model(input_var)
                            loss = criterion(output, target_var)
                            prec1 = accuracy(output.data, target, topk=(1,))[0]
                            eval_losses.update(loss.item(), input_img.size(0))
                            eval_top1.update(prec1.item(), input_img.size(0))