train_diffusion.py

import argparse
import wandb
import pkg_resources
import os
import torch.backends.cudnn as cudnn
import torch
import numpy as np
from torch import optim
from tqdm import tqdm
import wandb
from model.smpl.SMPL import SMPL_layer
from utils.tools import set_random_seed, get_config, print_args, create_directory_if_not_exists, find_arg_first_occurrence
from utils.learning import AverageMeter, save_model
from utils.learning_diffusion import load_data, load_model, batch_compute_similarity_transform_torch
from loss.diffusion_loss import DiffusionLoss
import utils.misc as misc
import configs.constants as _C
import torch.nn as nn
from utils.augmentation import Compose, TimeHalfMirrorning, TimeReverse, TimeWarping, TimePermutation, FeatureDrop



@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
                      for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-sd', '--seed', default=0,
                        type=int, help='random seed')
    parser.add_argument(
        "--config", type=str, default="configs/diffusion_twist.yaml", help="Path to the config file.")
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--eval-only', action='store_true')
    parser.add_argument('--final-test', action='store_true')
    parser.add_argument('-c', '--checkpoint', type=str, metavar='PATH', default='checkpoint',
                        help='checkpoint directory')
    parser.add_argument('--num-cpus', default=16, type=int,
                        help='Number of CPU cores')
    parser.add_argument('--use-wandb', action='store_true')
    parser.add_argument('--wandb-name', default=None, type=str)
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')

    # distributed training parameters
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--local_rank', default=-1, type=int)
    parser.add_argument('--dist_on_itp', action='store_true')
    parser.add_argument('--dist_url', default='env://',
                        help='url used to set up distributed training')
    opts = parser.parse_args()
    return opts


def train_one_epoch(args, augmentations, model, dataloader, optimizer, criterion, device):
    model.train()
    meters = {
        'loss_total': AverageMeter(),
        'loss_denoiser': AverageMeter(),
        'loss_score': AverageMeter(),
        'loss_pose': AverageMeter(),
        'loss_pose_vel': AverageMeter(),
        'loss_noise': AverageMeter(),
        'loss_vel': AverageMeter(),
        'loss_limb': AverageMeter(),
        'loss_pose_2d': AverageMeter(),
        'similarity_score': AverageMeter()
    }

    for item in tqdm(dataloader):
        for key in item:
            if not isinstance(item[key], list):
                item[key] = item[key].to(device, non_blocking=True)
                if item[key].dtype == torch.float64:
                    item[key] = item[key].float()
        frame_features = item['features']

        pose_skeleton = item['pose_skeleton']
        phis = item['phis']

        B, T = pose_skeleton.shape[:2]

        pose = torch.cat([pose_skeleton.reshape(B, T, -1),
                            phis.reshape(B, T, -1)], dim=-1)

        frame_features, pose = augmentations(frame_features, pose)

        out = model(frame_features, pose)

        optimizer.zero_grad()
        critertion_out = criterion(out, pose, item)
        critertion_out['loss_total'].backward()
        optimizer.step()

        for meter in meters:
            if meter in critertion_out:
                meters[meter].update(critertion_out[meter].item(), B)

    return {key: meter.avg for key, meter in meters.items() if meter.avg > 0}


def evaluate(args, model, dataloader, body_models, device, is_secondary=False):    
    model.eval()
    

    mpjpe_meter, mve_meter, pa_mpjpe_meter = AverageMeter(), AverageMeter(), AverageMeter()
    accel_meter = AverageMeter()
    
    dataset_name = args.main_validation_dataset if not is_secondary else args.secondary_validation_dataset

    is_joints24 = dataset_name in ['emdb', 'rich']
    is_3dpw = dataset_name == '3dpw'

    with torch.no_grad():
        for item in tqdm(dataloader):
            for key in item:
                if not isinstance(item[key], list):
                    item[key] = item[key].to(device, non_blocking=True)
                    if item[key].dtype == torch.float64:
                        item[key] = item[key].float()
            frame_features = item['features']
            flipped_features = item['flipped_features']
            pose_skeleton = item['pose_skeleton']
            phis = item['phis']
            betas = item['betas']
            hmr_betas = item[f'{args.backbone_type}_betas']
            frame_ids = item['frame_ids']
            gender = item['gender'][0]
            if gender not in ['male', 'female']:
                gender = {
                    0: 'male',
                    1: 'female'
                }[gender.item()]
            

            # Some of the clips less than 243 frames are resampled to become 243 frames
            # batch_mask helps us consider them only once in the evaluation.
            arg_first = find_arg_first_occurrence(frame_ids)
            batch_mask = torch.zeros_like(arg_first, dtype=bool)
            for i in range(batch_mask.shape[0]):
                batch_mask[i, arg_first[i][arg_first[i] >= 0]] = 1

            out_pred_joints_multi_hypo, out_pred_vertices_multi_hypo = [], []
            for _ in range(1):
                reconstructed_pose = model(frame_features)
                if not (flipped_features == 0).all():
                    reconstructed_pose_flipped = model(flipped_features)

                    # Flip back the flipped result and refine the pose skeleton by taking the average
                    skeleton_pred = reconstructed_pose[...,
                                                        :29*3].reshape(-1, 29, 3)
                    skeleton_pred_flipped = reconstructed_pose_flipped[...,
                                                                        :29*3].reshape(-1, 29, 3)
                    left_joints = [10, 7, 4, 1, 13, 16, 18, 20, 22, 25, 27]
                    right_joints = [11, 8, 5, 2,
                                    14, 17, 19, 21, 23, 26, 28]
                    skeleton_pred_flipped[..., 0] *= -1
                    skeleton_pred_flipped[..., left_joints + right_joints,
                                            :] = skeleton_pred_flipped[..., right_joints + left_joints, :]
                    refined_skeleton_pred = (
                        skeleton_pred + skeleton_pred_flipped) / 2
                    reconstructed_pose[..., :29 *
                                        3] = refined_skeleton_pred.reshape(-1, args.n_frames, 29*3)

                batch_size = pose_skeleton.shape[0]

                hmr_betas = hmr_betas.reshape(batch_size, -1, 10).mean(
                    dim=1, keepdims=True).repeat(1, args.n_frames, 1)
                
                out_pred = body_models['neutral'].hybrik(
                    pose_skeleton=reconstructed_pose[...,
                                                     :29*3].reshape(args.n_frames * batch_size, 29, 3),
                    phis=reconstructed_pose[..., 29 *
                                            3:].reshape(args.n_frames * batch_size, 23, 2),
                    betas=hmr_betas.reshape(args.n_frames * batch_size, 10),
                    global_orient=None
                )
                out_pred_vertices = out_pred.vertices.reshape(
                    batch_size, args.n_frames, -1, 3)
                if is_joints24:
                    out_pred_joints = out_pred.joints.reshape(
                        batch_size, args.n_frames, -1, 3)
                else:
                    out_pred_joints = out_pred.joints_from_verts.reshape(
                        batch_size, args.n_frames, -1, 3)

                pred_pelvis = out_pred_joints[:, :, 0:1]
                out_pred_joints = out_pred_joints - pred_pelvis
                out_pred_vertices = out_pred_vertices - pred_pelvis
                if is_3dpw:
                    out_pred_joints = out_pred_joints[:, :, _C.KEYPOINTS.H36M_TO_J14]
                out_pred_joints_multi_hypo.append(out_pred_joints.unsqueeze(1))
                out_pred_vertices_multi_hypo.append(out_pred_vertices.unsqueeze(1))
            out_pred_joints = torch.cat(out_pred_joints_multi_hypo, dim=1)
            out_pred_vertices = torch.cat(out_pred_vertices_multi_hypo, dim=1)

            # out_pred_joints = torch.stack(out_pred_joints_multi_hypo).mean(dim=0)
            # out_pred_vertices = torch.stack(out_pred_vertices_multi_hypo).mean(dim=0)

            out_gt = body_models[gender].hybrik(
                pose_skeleton=pose_skeleton.reshape(batch_size * args.n_frames, 29, 3),
                phis=phis.reshape(batch_size* args.n_frames, 23, 2),
                betas=betas.reshape(batch_size * args.n_frames, -1),
                global_orient=None
            )
            out_gts_vertices = out_gt.vertices.reshape(
                batch_size, args.n_frames, -1, 3)
            
            if is_joints24:
                out_gt_joints_smpl = out_gt.joints.reshape(
                    batch_size, args.n_frames, -1, 3)
            else:
                out_gt_joints_smpl = out_gt.joints_from_verts.reshape(
                    batch_size, args.n_frames, -1, 3)
            out_gts_joints = out_gt_joints_smpl
            gt_pelvis_smpl = out_gt_joints_smpl[:, :, 0:1]
            gt_pelvis = gt_pelvis_smpl

            out_gts_joints = (out_gts_joints - gt_pelvis)
            if is_3dpw:
                out_gts_joints = out_gts_joints[:, :, _C.KEYPOINTS.H36M_TO_J14]
            out_gts_vertices = out_gts_vertices - gt_pelvis_smpl

            mpjpe_hypos = torch.mean(torch.norm(
                out_pred_joints - out_gts_joints.unsqueeze(1), dim=-1), dim=(2, 3))

            min_idx = torch.argmin(mpjpe_hypos, dim=1)
            out_pred_joints = out_pred_joints[torch.arange(batch_size), min_idx]
            out_pred_vertices = out_pred_vertices[torch.arange(batch_size), min_idx]

            
            batch_mask = concat_all_gather(batch_mask).flatten()
            
            out_pred_joints = concat_all_gather(out_pred_joints)
            out_pred_joints = out_pred_joints.reshape(-1, 14 if is_3dpw else 24, 3)[batch_mask]
            
            out_gts_joints = concat_all_gather(out_gts_joints)
            out_gts_joints = out_gts_joints.reshape(-1, 14 if is_3dpw else 24, 3)[batch_mask]
            
            out_pred_vertices = concat_all_gather(out_pred_vertices)
            out_pred_vertices = out_pred_vertices.reshape(-1, 6890, 3)[batch_mask]
            
            out_gts_vertices = concat_all_gather(out_gts_vertices)
            out_gts_vertices = out_gts_vertices.reshape(-1, 6890, 3)[batch_mask]

            mpjpe = torch.mean(torch.norm(
                out_pred_joints - out_gts_joints, dim=-1)) * 1000
            
            mpjpe_meter.update(mpjpe.item(), out_pred_joints.shape[0])

            mve = torch.mean(torch.norm(out_pred_vertices -
                             out_gts_vertices, dim=-1)) * 1000
            mve_meter.update(mve.item(), out_pred_vertices.shape[0])

            S1_hat = batch_compute_similarity_transform_torch(
                out_pred_joints, out_gts_joints)
            pa_mpjpe = torch.mean(torch.norm(
                S1_hat - out_gts_joints, dim=-1)) * 1000
            pa_mpjpe_meter.update(pa_mpjpe.item(), out_pred_joints.shape[0])
            
            accel_gt = out_gts_joints[:-2] -  2 * out_gts_joints[1:-1] + out_gts_joints[2:]
            accel_pred = out_pred_joints[:-2] - 2 * out_pred_joints[1:-1] + out_pred_joints[2:]
            accel_err = torch.mean(torch.norm(accel_pred - accel_gt, dim=-1), dim=1)[1:-1]
            accel_err = accel_err * (30 ** 2)       # per frame^s to per s^2
            if accel_err.shape[0] == 0:
                continue
            accel_meter.update(torch.mean(accel_err).item(), out_pred_joints.shape[0])
            
            
                         
    if dataset_name == 'rich':
        data_str = '_rich'
    elif dataset_name == 'emdb':
        data_str = '_emdb'
    else:
        data_str = ""
        
    return {
        f'eval{data_str}/MPJPE': mpjpe_meter.avg,
        f'eval{data_str}/MVE': mve_meter.avg,
        f'eval{data_str}/PA-MPJPE': pa_mpjpe_meter.avg,
        f'eval{data_str}/ACCEL': accel_meter.avg
    }


def calculate_weight_magnitudes(model):
    weight_magnitudes = {}
    for name, param in model.named_parameters():
        if param.requires_grad:
            offset = 1 if name.split('.')[0] == 'module' else 0
            is_weight = 'weight' in name
            is_norm = 'norm' in name
            if not is_weight or is_norm:
                continue
            block_name = name.split(
                '.')[2 + offset] if 'encoder' in name else name.split('.')[1 + offset]
            if block_name in weight_magnitudes:
                weight_magnitudes[block_name] += param.data.norm().item()
            else:
                weight_magnitudes[block_name] = param.data.norm().item()
    weight_magnitudes = {'weight_magnitude/' +
                         key: value for key, value in weight_magnitudes.items()}
    return weight_magnitudes


def train(args, opts):
    print_args(args)
    if misc.is_main_process():
        create_directory_if_not_exists(opts.checkpoint)

    device = torch.device(opts.device)
    train_dataloader, validation_dataloader, secondary_validation_dataloader = load_data(args, opts, device)

    augmentations = Compose([
        TimeReverse(p=args.time_reverse_prob),
        TimePermutation(p=args.time_permutation_prob),
        TimeHalfMirrorning(p=args.time_half_mirrorning_prob),
        TimeWarping(p=args.time_warping_prob),
        FeatureDrop(p=args.feature_drop_prob, r=args.feature_drop_ratio)
    ])

    model, model_without_ddp = load_model(args, opts, device)
    J_regressor_h36m = torch.from_numpy(
            np.load(_C.BMODEL.JOINTS_REGRESSOR_H36M)
        )
    

    genders = ['neutral', 'male', 'female']
    body_models = {
        gender: SMPL_layer(model_path=os.path.join(_C.BMODEL.SMPL, f'SMPL_{gender.upper()}.pkl'),
                           h36m_jregressor=J_regressor_h36m).to(device)
        for gender in genders
    }
    for body_model in body_models.values():
        body_model.eval()
        for param in body_model.parameters():
            param.requires_grad = False

    optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model_without_ddp.parameters()),
                            lr=args.lr)
    criterion = DiffusionLoss(lambda_denoiser=args.lambda_denoiser,
                              lambda_score=args.lambda_score,
                              lambda_pose=args.lambda_pose,
                              lambda_limb=args.lambda_limb,
                              lambda_pose_vel=args.lambda_pose_vel,
                              lambda_hybrid1=args.lambda_hybrid1,
                              lambda_hybrid2=args.lambda_hybrid2,
                              lambda_2d_pose=args.lambda_2d_pose,
                              body_model=body_models['neutral']).to(device)

    last_ckpt_path = os.path.join(opts.checkpoint, 'last_ckpt.pth.tr')
    best_ckpt_path = os.path.join(opts.checkpoint, 'best_ckpt.pth.tr')
    resume = os.path.exists(last_ckpt_path)
    if resume:
        print('[INFO] Resuming from last checkpoint')
        checkpoint = torch.load(
            last_ckpt_path, map_location=lambda storage, loc: storage)
        model_without_ddp.load_state_dict(checkpoint['model'], strict=True)
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch'] + 1
        min_error = checkpoint['min_error']
        wandb_id = checkpoint['wandb_id']
    else:
        start_epoch = 0
        wandb_id = wandb.util.generate_id
        min_error = float('inf')

    if opts.use_wandb and misc.is_main_process():
        if resume:
            wandb.init(id=wandb_id,
                       project='VideoDiffHMR',
                       resume='must',
                       settings=wandb.Settings(start_method='fork'))
        else:
            wandb.init(name=opts.wandb_name,
                       project='VideoDiffHMR',
                       settings=wandb.Settings(start_method='fork'))
            wandb.config.update({"run_id": wandb_id})
            wandb.config.update(args)
            installed_packages = {
                d.project_name: d.version for d in pkg_resources.working_set}
            wandb.config.update({'installed_packages': installed_packages})
            wandb_id = wandb.run.id
            wandb.define_metric('eval/MPJPE', summary='min')
            wandb.define_metric('eval_emdb/MPJPE', summary='min')

    print(f'[INFO] Starting from epoch {start_epoch}')
    if opts.eval_only:
        try:
            validation_dataloader.sampler.set_epoch(start_epoch)
            secondary_validation_dataloader.sampler.set_epoch(start_epoch)
        except AttributeError:  # not running in distributed mode
            pass
        print(f"[INFO] Evaluation on {args.main_validation_dataset} dataset")
        eval_result = evaluate(
                args, model, validation_dataloader, body_models, device, is_secondary=False)
        print(eval_result)
        
        print(f"[INFO] Evaluation on {args.secondary_validation_dataset} dataset")
        eval_result = evaluate(
                args, model, secondary_validation_dataloader, body_models, device, is_secondary=True)
        print(eval_result)
        exit()

    for epoch in range(start_epoch, args.epochs):
        print(f"[INFO] Epoch {epoch}")
        if opts.distributed:
            train_dataloader.sampler.set_epoch(epoch)
            validation_dataloader.sampler.set_epoch(epoch)

        train_loss = train_one_epoch(
            args, augmentations, model, train_dataloader, optimizer, criterion, device)
        if args.evaluate_freq > 0 and epoch % args.evaluate_freq == 0 or epoch == args.epochs - 1:
            print("[INFO] Evaluation")
            eval_result = evaluate(
                args, model, validation_dataloader, body_models, device, is_secondary=False)
            if secondary_validation_dataloader is not None:
                eval_result_secondary = evaluate(
                    args, model, secondary_validation_dataloader, body_models, device, is_secondary=True)
            else:
                eval_result_secondary = {}
        else:
            eval_result = {}
            eval_result_secondary = {}
        
        eval_mpjpe = float('inf')
        for key in eval_result:
            if '/MPJPE' in key:
                eval_mpjpe = eval_result[key]
                break
        if eval_mpjpe < min_error:
            min_error = eval_mpjpe
            save_model(model_without_ddp, optimizer, min_error,
                       epoch, wandb_id, best_ckpt_path)
        if misc.is_main_process():
            weight_magnitudes = calculate_weight_magnitudes(model)
            log_dict = {
                **train_loss,
                **eval_result,
                'eval/min_MPJPE': min_error,
                **weight_magnitudes,
                **eval_result_secondary
            }
            if opts.use_wandb:
                wandb.log(log_dict, step=epoch + 1)

            for key, value in log_dict.items():
                print(f"[INFO] {key}: {value}")

        if epoch % args.save_freq == 0 and args.save_freq != -1 or epoch == args.epochs - 1:
            save_model(model_without_ddp, optimizer, min_error,
                       epoch, wandb_id, last_ckpt_path)

        train_dataloader.dataset.prepare_video_batch()


def main():
    opts = parse_args()

    misc.init_distributed_mode(opts)
    set_random_seed(opts.seed + misc.get_rank())
    cudnn.benchmark = True
    

    args = get_config(opts.config)

    if not opts.distributed:  # to work as identity function on single GPU
        global concat_all_gather
        def concat_all_gather(x): return x

    if opts.final_test:
        args.validation_dataset_path = args.validation_dataset_path.replace(
            "val", "test")

    train(args, opts)


if __name__ == "__main__":
    main()