mingyuan/MotionDiffuse

1import os2from os.path import join as pjoin3 4import utils.paramUtil as paramUtil5from options.train_options import TrainCompOptions6from utils.plot_script import *7 8from models import MotionTransformer9from trainers import DDPMTrainer10from datasets import Text2MotionDataset11 12from mmcv.runner import get_dist_info, init_dist13from mmcv.parallel import MMDistributedDataParallel14import torch15import torch.distributed as dist16 17 18def build_models(opt, dim_pose):19    encoder = MotionTransformer(20        input_feats=dim_pose,21        num_frames=opt.max_motion_length,22        num_layers=opt.num_layers,23        latent_dim=opt.latent_dim,24        no_clip=opt.no_clip,25        no_eff=opt.no_eff)26    return encoder27 28 29if __name__ == '__main__':30    parser = TrainCompOptions()31    opt = parser.parse()32    rank, world_size = get_dist_info()33 34    opt.device = torch.device("cuda")35    torch.autograd.set_detect_anomaly(True)36 37    opt.save_root = pjoin(opt.checkpoints_dir, opt.dataset_name, opt.name)38    opt.model_dir = pjoin(opt.save_root, 'model')39    opt.meta_dir = pjoin(opt.save_root, 'meta')40 41    if rank == 0:42        os.makedirs(opt.model_dir, exist_ok=True)43        os.makedirs(opt.meta_dir, exist_ok=True)44    if world_size > 1:45        dist.barrier()46 47    if opt.dataset_name == 't2m':48        opt.data_root = './data/HumanML3D'49        opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs')50        opt.text_dir = pjoin(opt.data_root, 'texts')51        opt.joints_num = 2252        radius = 453        fps = 2054        opt.max_motion_length = 19655        dim_pose = 26356        kinematic_chain = paramUtil.t2m_kinematic_chain57    elif opt.dataset_name == 'kit':58        opt.data_root = './data/KIT-ML'59        opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs')60        opt.text_dir = pjoin(opt.data_root, 'texts')61        opt.joints_num = 2162        radius = 240 * 863        fps = 12.564        dim_pose = 25165        opt.max_motion_length = 19666        kinematic_chain = paramUtil.kit_kinematic_chain67 68    else:69        raise KeyError('Dataset Does Not Exist')70 71    dim_word = 30072    mean = np.load(pjoin(opt.data_root, 'Mean.npy'))73    std = np.load(pjoin(opt.data_root, 'Std.npy'))74 75    train_split_file = pjoin(opt.data_root, 'train.txt')76 77    encoder = build_models(opt, dim_pose)78    if world_size > 1:79        encoder = MMDistributedDataParallel(80            encoder.cuda(),81            device_ids=[torch.cuda.current_device()],82            broadcast_buffers=False,83            find_unused_parameters=True)84    else:85        encoder = encoder.cuda()86 87    trainer = DDPMTrainer(opt, encoder)88    train_dataset = Text2MotionDataset(opt, mean, std, train_split_file, opt.times)89    trainer.train(train_dataset)90