import os import numpy as np import argparse import errno import math import pickle import tensorboardX from tqdm import tqdm from time import time import copy import random import prettytable import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader from lib.utils.tools import * from lib.utils.learning import * from lib.utils.utils_data import flip_data from lib.data.dataset_motion_2d import PoseTrackDataset2D, InstaVDataset2D from lib.data.dataset_motion_3d import MotionDataset3D from lib.data.augmentation import Augmenter2D from lib.data.datareader_h36m import DataReaderH36M from lib.model.loss import * def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, default="configs/pretrain.yaml", help="Path to the config file.") parser.add_argument('-c', '--checkpoint', default='checkpoint', type=str, metavar='PATH', help='checkpoint directory') parser.add_argument('-p', '--pretrained', default='checkpoint', type=str, metavar='PATH', help='pretrained checkpoint directory') parser.add_argument('-r', '--resume', default='', type=str, metavar='FILENAME', help='checkpoint to resume (file name)') parser.add_argument('-e', '--evaluate', default='', type=str, metavar='FILENAME', help='checkpoint to evaluate (file name)') parser.add_argument('-ms', '--selection', default='latest_epoch.bin', type=str, metavar='FILENAME', help='checkpoint to finetune (file name)') parser.add_argument('-sd', '--seed', default=0, type=int, help='random seed') opts = parser.parse_args() return opts def set_random_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) def save_checkpoint(chk_path, epoch, lr, optimizer, model_pos, min_loss): print('Saving checkpoint to', chk_path) torch.save({ 'epoch': epoch + 1, 'lr': lr, 'optimizer': optimizer.state_dict(), 'model_pos': model_pos.state_dict(), 'min_loss' : min_loss }, chk_path) def evaluate(args, model_pos, test_loader, datareader): print('INFO: Testing') results_all = [] model_pos.eval() with torch.no_grad(): for batch_input, batch_gt in tqdm(test_loader): N, T = batch_gt.shape[:2] if torch.cuda.is_available(): batch_input = batch_input.cuda() if args.no_conf: batch_input = batch_input[:, :, :, :2] if args.flip: batch_input_flip = flip_data(batch_input) predicted_3d_pos_1 = model_pos(batch_input) predicted_3d_pos_flip = model_pos(batch_input_flip) predicted_3d_pos_2 = flip_data(predicted_3d_pos_flip) # Flip back predicted_3d_pos = (predicted_3d_pos_1+predicted_3d_pos_2) / 2 else: predicted_3d_pos = model_pos(batch_input) if args.rootrel: predicted_3d_pos[:,:,0,:] = 0 # [N,T,17,3] else: batch_gt[:,0,0,2] = 0 if args.gt_2d: predicted_3d_pos[...,:2] = batch_input[...,:2] results_all.append(predicted_3d_pos.cpu().numpy()) results_all = np.concatenate(results_all) results_all = datareader.denormalize(results_all) _, split_id_test = datareader.get_split_id() actions = np.array(datareader.dt_dataset['test']['action']) factors = np.array(datareader.dt_dataset['test']['2.5d_factor']) gts = np.array(datareader.dt_dataset['test']['joints_2.5d_image']) sources = np.array(datareader.dt_dataset['test']['source']) num_test_frames = len(actions) frames = np.array(range(num_test_frames)) action_clips = actions[split_id_test] factor_clips = factors[split_id_test] source_clips = sources[split_id_test] frame_clips = frames[split_id_test] gt_clips = gts[split_id_test] assert len(results_all)==len(action_clips) e1_all = np.zeros(num_test_frames) e2_all = np.zeros(num_test_frames) oc = np.zeros(num_test_frames) results = {} results_procrustes = {} action_names = sorted(set(datareader.dt_dataset['test']['action'])) for action in action_names: results[action] = [] results_procrustes[action] = [] block_list = ['s_09_act_05_subact_02', 's_09_act_10_subact_02', 's_09_act_13_subact_01'] for idx in range(len(action_clips)): source = source_clips[idx][0][:-6] if source in block_list: continue frame_list = frame_clips[idx] action = action_clips[idx][0] factor = factor_clips[idx][:,None,None] gt = gt_clips[idx] pred = results_all[idx] pred *= factor # Root-relative Errors pred = pred - pred[:,0:1,:] gt = gt - gt[:,0:1,:] err1 = mpjpe(pred, gt) err2 = p_mpjpe(pred, gt) e1_all[frame_list] += err1 e2_all[frame_list] += err2 oc[frame_list] += 1 for idx in range(num_test_frames): if e1_all[idx] > 0: err1 = e1_all[idx] / oc[idx] err2 = e2_all[idx] / oc[idx] action = actions[idx] results[action].append(err1) results_procrustes[action].append(err2) final_result = [] final_result_procrustes = [] summary_table = prettytable.PrettyTable() summary_table.field_names = ['test_name'] + action_names for action in action_names: final_result.append(np.mean(results[action])) final_result_procrustes.append(np.mean(results_procrustes[action])) summary_table.add_row(['P1'] + final_result) summary_table.add_row(['P2'] + final_result_procrustes) print(summary_table) e1 = np.mean(np.array(final_result)) e2 = np.mean(np.array(final_result_procrustes)) print('Protocol #1 Error (MPJPE):', e1, 'mm') print('Protocol #2 Error (P-MPJPE):', e2, 'mm') print('----------') return e1, e2, results_all def train_epoch(args, model_pos, train_loader, losses, optimizer, has_3d, has_gt): model_pos.train() for idx, (batch_input, batch_gt) in tqdm(enumerate(train_loader)): batch_size = len(batch_input) if torch.cuda.is_available(): batch_input = batch_input.cuda() batch_gt = batch_gt.cuda() with torch.no_grad(): if args.no_conf: batch_input = batch_input[:, :, :, :2] if not has_3d: conf = copy.deepcopy(batch_input[:,:,:,2:]) # For 2D data, weight/confidence is at the last channel if args.rootrel: batch_gt = batch_gt - batch_gt[:,:,0:1,:] else: batch_gt[:,:,:,2] = batch_gt[:,:,:,2] - batch_gt[:,0:1,0:1,2] # Place the depth of first frame root to 0. if args.mask or args.noise: batch_input = args.aug.augment2D(batch_input, noise=(args.noise and has_gt), mask=args.mask) # Predict 3D poses predicted_3d_pos = model_pos(batch_input) # (N, T, 17, 3) optimizer.zero_grad() if has_3d: loss_3d_pos = loss_mpjpe(predicted_3d_pos, batch_gt) loss_3d_scale = n_mpjpe(predicted_3d_pos, batch_gt) loss_3d_velocity = loss_velocity(predicted_3d_pos, batch_gt) loss_lv = loss_limb_var(predicted_3d_pos) loss_lg = loss_limb_gt(predicted_3d_pos, batch_gt) loss_a = loss_angle(predicted_3d_pos, batch_gt) loss_av = loss_angle_velocity(predicted_3d_pos, batch_gt) loss_total = loss_3d_pos + \ args.lambda_scale * loss_3d_scale + \ args.lambda_3d_velocity * loss_3d_velocity + \ args.lambda_lv * loss_lv + \ args.lambda_lg * loss_lg + \ args.lambda_a * loss_a + \ args.lambda_av * loss_av losses['3d_pos'].update(loss_3d_pos.item(), batch_size) losses['3d_scale'].update(loss_3d_scale.item(), batch_size) losses['3d_velocity'].update(loss_3d_velocity.item(), batch_size) losses['lv'].update(loss_lv.item(), batch_size) losses['lg'].update(loss_lg.item(), batch_size) losses['angle'].update(loss_a.item(), batch_size) losses['angle_velocity'].update(loss_av.item(), batch_size) losses['total'].update(loss_total.item(), batch_size) else: loss_2d_proj = loss_2d_weighted(predicted_3d_pos, batch_gt, conf) loss_total = loss_2d_proj losses['2d_proj'].update(loss_2d_proj.item(), batch_size) losses['total'].update(loss_total.item(), batch_size) loss_total.backward() optimizer.step() def train_with_config(args, opts): print(args) try: os.makedirs(opts.checkpoint) except OSError as e: if e.errno != errno.EEXIST: raise RuntimeError('Unable to create checkpoint directory:', opts.checkpoint) train_writer = tensorboardX.SummaryWriter(os.path.join(opts.checkpoint, "logs")) print('Loading dataset...') trainloader_params = { 'batch_size': args.batch_size, 'shuffle': True, 'num_workers': 12, 'pin_memory': True, 'prefetch_factor': 4, 'persistent_workers': True } testloader_params = { 'batch_size': args.batch_size, 'shuffle': False, 'num_workers': 12, 'pin_memory': True, 'prefetch_factor': 4, 'persistent_workers': True } train_dataset = MotionDataset3D(args, args.subset_list, 'train') test_dataset = MotionDataset3D(args, args.subset_list, 'test') train_loader_3d = DataLoader(train_dataset, **trainloader_params) test_loader = DataLoader(test_dataset, **testloader_params) if args.train_2d: posetrack = PoseTrackDataset2D() posetrack_loader_2d = DataLoader(posetrack, **trainloader_params) instav = InstaVDataset2D() instav_loader_2d = DataLoader(instav, **trainloader_params) datareader = DataReaderH36M(n_frames=args.clip_len, sample_stride=args.sample_stride, data_stride_train=args.data_stride, data_stride_test=args.clip_len, dt_root = 'data/motion3d', dt_file=args.dt_file) min_loss = 100000 model_backbone = load_backbone(args) model_params = 0 for parameter in model_backbone.parameters(): model_params = model_params + parameter.numel() print('INFO: Trainable parameter count:', model_params) if torch.cuda.is_available(): model_backbone = nn.DataParallel(model_backbone) model_backbone = model_backbone.cuda() if args.finetune: if opts.resume or opts.evaluate: chk_filename = opts.evaluate if opts.evaluate else opts.resume print('Loading checkpoint', chk_filename) checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage) model_backbone.load_state_dict(checkpoint['model_pos'], strict=True) model_pos = model_backbone else: chk_filename = os.path.join(opts.pretrained, opts.selection) print('Loading checkpoint', chk_filename) checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage) model_backbone.load_state_dict(checkpoint['model_pos'], strict=True) model_pos = model_backbone else: chk_filename = os.path.join(opts.checkpoint, "latest_epoch.bin") if os.path.exists(chk_filename): opts.resume = chk_filename if opts.resume or opts.evaluate: chk_filename = opts.evaluate if opts.evaluate else opts.resume print('Loading checkpoint', chk_filename) checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage) model_backbone.load_state_dict(checkpoint['model_pos'], strict=True) model_pos = model_backbone if args.partial_train: model_pos = partial_train_layers(model_pos, args.partial_train) if not opts.evaluate: lr = args.learning_rate optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model_pos.parameters()), lr=lr, weight_decay=args.weight_decay) lr_decay = args.lr_decay st = 0 if args.train_2d: print('INFO: Training on {}(3D)+{}(2D) batches'.format(len(train_loader_3d), len(instav_loader_2d) + len(posetrack_loader_2d))) else: print('INFO: Training on {}(3D) batches'.format(len(train_loader_3d))) if opts.resume: st = checkpoint['epoch'] if 'optimizer' in checkpoint and checkpoint['optimizer'] is not None: optimizer.load_state_dict(checkpoint['optimizer']) else: print('WARNING: this checkpoint does not contain an optimizer state. The optimizer will be reinitialized.') lr = checkpoint['lr'] if 'min_loss' in checkpoint and checkpoint['min_loss'] is not None: min_loss = checkpoint['min_loss'] args.mask = (args.mask_ratio > 0 and args.mask_T_ratio > 0) if args.mask or args.noise: args.aug = Augmenter2D(args) # Training for epoch in range(st, args.epochs): print('Training epoch %d.' % epoch) start_time = time() losses = {} losses['3d_pos'] = AverageMeter() losses['3d_scale'] = AverageMeter() losses['2d_proj'] = AverageMeter() losses['lg'] = AverageMeter() losses['lv'] = AverageMeter() losses['total'] = AverageMeter() losses['3d_velocity'] = AverageMeter() losses['angle'] = AverageMeter() losses['angle_velocity'] = AverageMeter() N = 0 # Curriculum Learning if args.train_2d and (epoch >= args.pretrain_3d_curriculum): train_epoch(args, model_pos, posetrack_loader_2d, losses, optimizer, has_3d=False, has_gt=True) train_epoch(args, model_pos, instav_loader_2d, losses, optimizer, has_3d=False, has_gt=False) train_epoch(args, model_pos, train_loader_3d, losses, optimizer, has_3d=True, has_gt=True) elapsed = (time() - start_time) / 60 if args.no_eval: print('[%d] time %.2f lr %f 3d_train %f' % ( epoch + 1, elapsed, lr, losses['3d_pos'].avg)) else: e1, e2, results_all = evaluate(args, model_pos, test_loader, datareader) print('[%d] time %.2f lr %f 3d_train %f e1 %f e2 %f' % ( epoch + 1, elapsed, lr, losses['3d_pos'].avg, e1, e2)) train_writer.add_scalar('Error P1', e1, epoch + 1) train_writer.add_scalar('Error P2', e2, epoch + 1) train_writer.add_scalar('loss_3d_pos', losses['3d_pos'].avg, epoch + 1) train_writer.add_scalar('loss_2d_proj', losses['2d_proj'].avg, epoch + 1) train_writer.add_scalar('loss_3d_scale', losses['3d_scale'].avg, epoch + 1) train_writer.add_scalar('loss_3d_velocity', losses['3d_velocity'].avg, epoch + 1) train_writer.add_scalar('loss_lv', losses['lv'].avg, epoch + 1) train_writer.add_scalar('loss_lg', losses['lg'].avg, epoch + 1) train_writer.add_scalar('loss_a', losses['angle'].avg, epoch + 1) train_writer.add_scalar('loss_av', losses['angle_velocity'].avg, epoch + 1) train_writer.add_scalar('loss_total', losses['total'].avg, epoch + 1) # Decay learning rate exponentially lr *= lr_decay for param_group in optimizer.param_groups: param_group['lr'] *= lr_decay # Save checkpoints chk_path = os.path.join(opts.checkpoint, 'epoch_{}.bin'.format(epoch)) chk_path_latest = os.path.join(opts.checkpoint, 'latest_epoch.bin') chk_path_best = os.path.join(opts.checkpoint, 'best_epoch.bin'.format(epoch)) save_checkpoint(chk_path_latest, epoch, lr, optimizer, model_pos, min_loss) if (epoch + 1) % args.checkpoint_frequency == 0: save_checkpoint(chk_path, epoch, lr, optimizer, model_pos, min_loss) if e1 < min_loss: min_loss = e1 save_checkpoint(chk_path_best, epoch, lr, optimizer, model_pos, min_loss) if opts.evaluate: e1, e2, results_all = evaluate(args, model_pos, test_loader, datareader) if __name__ == "__main__": opts = parse_args() set_random_seed(opts.seed) args = get_config(opts.config) train_with_config(args, opts)