Ejemplo n.º 1
0
def step(split, epoch, opt, dataLoader, model, criterion, optimizer=None):
    if split == 'train':
        model.train()
    else:
        model.eval()
    Loss, Acc, Mpjpe, Loss3D = AverageMeter(), AverageMeter(), AverageMeter(
    ), AverageMeter()

    nIters = len(dataLoader)
    bar = Bar('==>', max=nIters)

    for i, (input, target2D, target3D, meta) in enumerate(dataLoader):
        input_var = torch.autograd.Variable(input).float().cuda()
        target2D_var = torch.autograd.Variable(target2D).float().cuda()
        target3D_var = torch.autograd.Variable(target3D).float().cuda()

        output = model(input_var)
        reg = output[opt.nStack]
        if opt.DEBUG >= 2:
            gt = getPreds(target2D.cpu().numpy()) * 4
            pred = getPreds((output[opt.nStack - 1].data).cpu().numpy()) * 4
            debugger = Debugger()
            debugger.addImg(
                (input[0].numpy().transpose(1, 2, 0) * 256).astype(np.uint8))
            debugger.addPoint2D(pred[0], (255, 0, 0))
            debugger.addPoint2D(gt[0], (0, 0, 255))
            debugger.showImg()
            debugger.saveImg('debug/{}.png'.format(i))

        loss = FusionCriterion(opt.regWeight, opt.varWeight)(reg, target3D_var)
        Loss3D.update(loss.data[0], input.size(0))
        for k in range(opt.nStack):
            loss += criterion(output[k], target2D_var)

        Loss.update(loss.data[0], input.size(0))
        Acc.update(
            Accuracy((output[opt.nStack - 1].data).cpu().numpy(),
                     (target2D_var.data).cpu().numpy()))
        mpjpe, num3D = MPJPE((output[opt.nStack - 1].data).cpu().numpy(),
                             (reg.data).cpu().numpy(), meta)
        if num3D > 0:
            Mpjpe.update(mpjpe, num3D)
        if split == 'train':
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        Bar.suffix = '{split} Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | Loss3D {loss3d.avg:.6f} | Acc {Acc.avg:.6f} | Mpjpe {Mpjpe.avg:.6f} ({Mpjpe.val:.6f})'.format(
            epoch,
            i,
            nIters,
            total=bar.elapsed_td,
            eta=bar.eta_td,
            loss=Loss,
            Acc=Acc,
            split=split,
            Mpjpe=Mpjpe,
            loss3d=Loss3D)
        bar.next()

    bar.finish()
    return Loss.avg, Acc.avg, Mpjpe.avg, Loss3D.avg
Ejemplo n.º 2
0
def train(epoch):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()
    Loss, Acc, Mpjpe, Loss3D = AverageMeter(), AverageMeter(), AverageMeter(
    ), AverageMeter()

    # switch to train mode
    model.train()

    end = time.time()
    for i, data in enumerate(train_loader):
        input, target2D, target3D, meta = data
        input_var = torch.autograd.Variable(input).float().cuda()
        target2D_var = torch.autograd.Variable(target2D).float().cuda()
        target3D_var = torch.autograd.Variable(target3D).float().cuda()

        # 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)
        output = model(input_var)
        reg = output[args.nStack]

        # measure accuracy and record loss
        #prec1, prec5 = accuracy(output, target, topk=(1, 5))
        #losses.update(loss.item(), input.size(0))
        #top1.update(prec1[0], input.size(0))
        #top5.update(prec5[0], input.size(0))
        #print(reg)
        #print(reg.float())
        #print(reg.type(torch.FloatTensor))
        reg = reg.float().cuda()
        #print(reg.type())
        #loss = FusionCriterion(args.regWeight, args.varWeight)(reg, target3D_var)
        #Loss3D.update(loss.data[0], input.size(0))
        loss = 0
        for k in range(args.nStack):
            loss += criterion(output[k], target2D_var)

        Loss.update(loss.data[0], input.size(0))
        Acc.update(
            Accuracy((output[args.nStack - 1].data).cpu().numpy(),
                     (target2D_var.data).cpu().numpy()))
        mpjpe, num3D = MPJPE((output[args.nStack - 1].data).cpu().numpy(),
                             (reg.data).cpu().numpy(), meta)
        if num3D > 0:
            Mpjpe.update(mpjpe, num3D)

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

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

        loss_train, acc_train, mpjpe_train, loss3d_train = Loss.avg, Acc.avg, Mpjpe.avg, Loss3D.avg

        if i % args.print_freq == 0:
            print('Epoch: [{0}][{1}/{2}]\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})\t'
                  'Acc {acc.val:.3f} ({acc.avg:.3f})\t'
                  'Mpjpe {mpjpe.val:.3f} ({mpjpe.avg:.3f})\t'
                  'Loss3d {loss3d.val:.3f} ({loss3d.avg:.3f})\t'.format(
                      epoch,
                      i,
                      len(train_loader),
                      batch_time=batch_time,
                      data_time=data_time,
                      loss=Loss,
                      acc=Acc,
                      mpjpe=Mpjpe,
                      loss3d=Loss3D))
Ejemplo n.º 3
0
        output = model(input_var)

        reg = output[-1]

        loss = criterion(reg, target3D_var)
        Loss3D.update(loss.data[0], input1.size(0))

        print(len(output))
        for k in range(0, 3):
            loss += criterion(output[k], target2D_var)

        Loss.update(loss.data[0], input1.size(0))
        Acc.update(
            Accuracy((output[-2].data).cpu().numpy(),
                     (target2D_var.data).numpy()))
        mpjpe, num3D = MPJPE((output[-2].data).cpu().numpy(),
                             (reg[:, 32:].data).numpy(), meta)
        if num3D > 0:
            Mpjpe.update(mpjpe, num3D)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        Bar.suffix = 'Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | Loss3D {loss3d.avg:.6f} | Acc {Acc.avg:.6f} | Mpjpe {Mpjpe.avg:.6f}'.format(
            epoch,
            i,
            nIters,
            total=bar.elapsed_td,
            eta=bar.eta_td,
            loss=Loss,
            Acc=Acc,
            loss3d=Loss3D,
Ejemplo n.º 4
0
        #print(target2D,targets_3D)
        output = model(inputs)

        reg = output[opt.nStack]

        loss = FusionCriterion(opt.regWeight, opt.varWeight)(reg, targets_3D)
        Loss3D.update(loss.item(), input.size(0))

        for j in range(1, opt.nStack):
            loss += criterion(output[j], targets_2D)

        Loss.update(loss.item(), input.size(0))
        acc = Accuracy((output[opt.nStack - 1].data).cpu().numpy(),
                       (targets_2D.data).cpu().numpy())
        Acc.update(acc)
        mpjpe, num3D = MPJPE((output[opt.nStack - 1].data).cpu().numpy(),
                             (reg.data).cpu().numpy(), meta)
        if num3D > 0:
            Mpjpe.update(mpjpe, num3D)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print(epoch, i, loss.data, acc)

    if epoch % opt.valIntervals == 0:
        torch.save(model,
                   os.path.join(opt.saveDir, 'model_{}.pth'.format(epoch)))
    adjust_learning_rate(optimizer, epoch, opt.dropLR, opt.LR)

    print('Loss:', Loss.avg, 'Acc:', Acc.avg, 'Mpjpe:', Mpjpe.avg, 'Loss3D',
          Loss3D.avg)