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
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))
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,
#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)