train_ACMDM_Mesh.py

import os
from os.path import join as pjoin
import torch
import numpy as np
import random
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import torch.optim as optim
from models.ACMDM import ACMDM_models
from utils.datasets import Text2MotionDataset
import time
import copy
from collections import OrderedDict, defaultdict
from utils.train_utils import update_lr_warm_up, def_value, save, print_current_loss, update_ema
import argparse


def main(args):
    #################################################################################
    #                                      Seed                                     #
    #################################################################################
    torch.backends.cudnn.benchmark = False
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.autograd.set_detect_anomaly(True)
    # setting this to true significantly increase training and sampling speed
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    #################################################################################
    #                                    Train Data                                 #
    #################################################################################
    if args.dataset_name == "t2m":
        data_root = f'{args.dataset_dir}/HumanML3D/'
        dim_pose = 12
    else:
        raise NotImplementedError
    motion_dir = pjoin(data_root, 'meshes_after_ae') # use preprocessed for faster training
    text_dir = pjoin(data_root, 'texts')
    mean = np.load(pjoin(args.checkpoints_dir, args.dataset_name, args.ae_name, args.after_mean)) # make sure this is either calculated or obtained from given checkpoint files
    std = np.load(pjoin(args.checkpoints_dir, args.dataset_name, args.ae_name, args.after_std)) # make sure this is either calculated or obtained from given checkpoint files
    train_split_file = pjoin(data_root, 'train.txt')
    val_split_file = pjoin(data_root, 'val.txt')

    train_dataset = Text2MotionDataset(mean, std, train_split_file, args.dataset_name, motion_dir, text_dir,
                                          args.unit_length, args.max_motion_length, 20, evaluation=False, is_mesh=True)
    val_dataset = Text2MotionDataset(mean, std, val_split_file, args.dataset_name, motion_dir, text_dir,
                                          args.unit_length, args.max_motion_length, 20, evaluation=False, is_mesh=True)

    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, drop_last=True, num_workers=args.num_workers,
                              shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, drop_last=True, num_workers=args.num_workers,
                            shuffle=True)
    #################################################################################
    #                                      Models                                   #
    #################################################################################
    model_dir = pjoin(args.checkpoints_dir, args.dataset_name, args.name, 'model')
    os.makedirs(model_dir, exist_ok=True)

    acmdm = ACMDM_models[args.model](input_dim=dim_pose, cond_mode='text')
    ema_acmdm = copy.deepcopy(acmdm)
    ema_acmdm.eval()
    for param in ema_acmdm.parameters():
        param.requires_grad_(False)

    all_params = 0
    pc_transformer = sum(param.numel() for param in [p for name, p in acmdm.named_parameters() if not name.startswith('clip_model.')])
    all_params += pc_transformer
    print('Total parameters of all models: {:.2f}M'.format(all_params / 1000_000))

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    #################################################################################
    #                                    Training Loop                              #
    #################################################################################
    logger = SummaryWriter(model_dir)
    acmdm.to(device)
    ema_acmdm.to(device)

    optimizer = optim.AdamW(acmdm.parameters(), betas=(0.9, 0.99), lr=args.lr, weight_decay=1e-5)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=args.lr_decay)

    epoch = 0
    it = 0
    if args.is_continue:
        checkpoint = torch.load(pjoin(model_dir, 'latest.tar'), map_location=device)
        missing_keys, unexpected_keys = acmdm.load_state_dict(checkpoint['acmdm'], strict=False)
        missing_keys2, unexpected_keys2 = ema_acmdm.load_state_dict(checkpoint['ema_acmdm'], strict=False)
        assert len(unexpected_keys) == 0
        assert len(unexpected_keys2) == 0
        assert all([k.startswith('clip_model.') for k in missing_keys])
        assert all([k.startswith('clip_model.') for k in missing_keys2])
        optimizer.load_state_dict(checkpoint['opt_acmdm'])
        scheduler.load_state_dict(checkpoint['scheduler'])
        epoch, it = checkpoint['ep']+1, checkpoint['total_it']
        print("Load model epoch:%d iterations:%d" % (epoch, it))

    start_time = time.time()
    total_iters = args.epoch * len(train_loader)
    print(f'Total Epochs: {args.epoch}, Total Iters: {total_iters}')
    print('Iters Per Epoch, Training: %04d, Validation: %03d' % (len(train_loader), len(val_loader)))

    logs = defaultdict(def_value, OrderedDict())
    worst_loss = 100

    while epoch < args.epoch:
        acmdm.train()
        optimizer.zero_grad()

        for i, batch_data in enumerate(train_loader):
            it += 1
            if it < args.warm_up_iter:
                update_lr_warm_up(it, args.warm_up_iter, optimizer, args.lr)

            conds, motion, m_lens = batch_data
            motion = motion.detach().float().to(device)
            m_lens = m_lens.detach().long().to(device)

            latent = motion

            conds = conds.to(device).float() if torch.is_tensor(conds) else conds

            # a naive way of writing gradient accumulation
            accum = args.accum
            accum_bs = args.batch_size // accum

            for kk in range(accum):
                loss = acmdm.forward_loss(latent.permute(0, 3, 1, 2)[kk * accum_bs:(kk + 1) * accum_bs],
                                          conds[kk * accum_bs:(kk + 1) * accum_bs],
                                          m_lens[kk * accum_bs:(kk + 1) * accum_bs])
                loss = loss / accum
                loss.backward()
                logs['loss'] += loss.item()
                logs['lr'] += (optimizer.param_groups[0]['lr'])/accum

            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()

            update_ema(acmdm, ema_acmdm, 0.9999)

            if it % args.log_every == 0:
                mean_loss = OrderedDict()
                for tag, value in logs.items():
                    logger.add_scalar('Train/%s' % tag, value / args.log_every, it)
                    mean_loss[tag] = value / args.log_every
                logs = defaultdict(def_value, OrderedDict())
                print_current_loss(start_time, it, total_iters, mean_loss, epoch=epoch, inner_iter=i)

        save(pjoin(model_dir, 'latest.tar'), epoch, acmdm, optimizer, scheduler, it, 'acmdm', ema=ema_acmdm)
        #################################################################################
        #                                      Eval Loop                                #
        #################################################################################
        print('Validation time:')
        acmdm.eval()
        val_loss = []
        with torch.no_grad():
            for i, batch_data in enumerate(val_loader):
                conds, motion, m_lens = batch_data
                motion = motion.detach().float().to(device)
                m_lens = m_lens.detach().long().to(device)

                latent = motion

                conds = conds.to(device).float() if torch.is_tensor(conds) else conds

                loss = acmdm.forward_loss(latent.permute(0,3,1,2), conds, m_lens)
                val_loss.append(loss.item())

        print(f"Validation loss:{np.mean(val_loss):.3f}")
        logger.add_scalar('Val/loss', np.mean(val_loss), epoch)
        if np.mean(val_loss) < worst_loss:
            print(f"Improved loss from {worst_loss:.02f} to {np.mean(val_loss)}!!!")
            worst_loss = np.mean(val_loss)
        epoch += 1


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--name', type=str, default='ACMDM_Mesh_Flow_S_PatchSize22')
    parser.add_argument('--ae_name', type=str, default="AE_Mesh")
    parser.add_argument('--is_ae', action="store_true")
    parser.add_argument('--after_mean', type=str, default='AE_Mesh_Post_Mean.npy')
    parser.add_argument('--after_std', type=str, default='AE_Mesh_Post_Std.npy')
    parser.add_argument('--model', type=str, default='ACMDM-Mesh-Flow-S-PatchSize22')
    parser.add_argument('--dataset_name', type=str, default='t2m')
    parser.add_argument('--dataset_dir', type=str, default='./datasets')
    parser.add_argument("--max_motion_length", type=int, default=196)
    parser.add_argument("--unit_length", type=int, default=4)
    parser.add_argument('--batch_size', default=64, type=int)

    parser.add_argument('--epoch', default=500, type=int)
    parser.add_argument('--warm_up_iter', default=2000, type=int)
    parser.add_argument('--accum', default=1, type=int)
    parser.add_argument('--lr', default=2e-4, type=float)
    parser.add_argument('--milestones', default=[50_000], nargs="+", type=int)
    parser.add_argument('--lr_decay', default=0.1, type=float)

    parser.add_argument("--seed", type=int, default=3407)
    parser.add_argument("--num_workers", type=int, default=8)
    parser.add_argument('--is_continue', action="store_true")
    parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints')

    parser.add_argument('--log_every', default=50, type=int)

    arg = parser.parse_args()
    main(arg)