train_seq.py

import multiprocessing
import os
from zlib import Z_FULL_FLUSH
# os.environ["WANDB_API_KEY"] = "your WANDB_API_KEY" # 
# os.environ["WANDB_MODE"] = "online"
# os.environ["CUDA_VISIBLE_DEVICES"] = "4,5,6,7"
import pickle
from functools import partial
from pathlib import Path
from args import FineDance_parse_train_opt, save_arguments_to_yaml
import sys

import torch
import torch.nn.functional as F
import wandb
from accelerate import Accelerator, DistributedDataParallelKwargs
from accelerate.state import AcceleratorState
from torch.utils.data import DataLoader
from tqdm import tqdm

from dataset.FineDance_dataset import FineDance_Smpl
from dataset.preprocess import increment_path
from dataset.preprocess import My_Normalizer as Normalizer        # do not use Normalizer
from model.adan import Adan
from model.diffusion import GaussianDiffusion
from model.model import DanceDecoder, SeqModel
from vis import SMPLX_Skeleton, SMPLSkeleton


def wrap(x):
    return {f"module.{key}": value for key, value in x.items()}


def maybe_wrap(x, num):
    return x if num == 1 else wrap(x)


class EDGE:
    def __init__(
        self,
        opt,
        feature_type,
        checkpoint_path="",
        normalizer=None,
        EMA=True,
        learning_rate=4e-4,
        weight_decay=0.02,
    ):
        self.opt = opt
        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
        self.accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
        state = AcceleratorState()
        num_processes = state.num_processes
    
        self.repr_dim = repr_dim = opt.nfeats
        feature_dim = 35
   
        self.horizon = horizon = opt.full_seq_len

        self.accelerator.wait_for_everyone()

        self.resume_num = 0
        checkpoint = None
        self.normalizer = None
        if checkpoint_path != "":
            checkpoint = torch.load(
                checkpoint_path, map_location=self.accelerator.device
            )
            self.resume_num = int(os.path.basename(checkpoint_path).split("-")[1].split(".")[0])      # int(os.path.basenam

        model = SeqModel(
            nfeats=repr_dim,
            seq_len=horizon,
            latent_dim=512,
            ff_size=1024,
            num_layers=8,
            num_heads=8,
            dropout=0.1,
            cond_feature_dim=feature_dim,
            activation=F.gelu,
        )
        if opt.nfeats == 139 or opt.nfeats == 135:
            smplx_fk = SMPLSkeleton(device=self.accelerator.device)
        else:
            smplx_fk = SMPLX_Skeleton(device=self.accelerator.device, batch=512000)
        diffusion = GaussianDiffusion(
            model,
            opt,
            horizon,
            repr_dim,
            smplx_model = smplx_fk,
            schedule="cosine",
            n_timestep=1000,
            predict_epsilon=False,
            loss_type="l2",
            use_p2=False,
            cond_drop_prob=0.25,
            guidance_weight=2,
            do_normalize = opt.do_normalize
        )

        print(
            "Model has {} parameters".format(sum(y.numel() for y in model.parameters()))
        )

        self.model = self.accelerator.prepare(model)                                            
        self.diffusion = diffusion.to(self.accelerator.device)                              # 为什么这里不需要prepare
        self.smplx_fk = smplx_fk     # to(self.accelerator.device)
        optim = Adan(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
        self.optim = self.accelerator.prepare(optim)

        if checkpoint_path != "":
            self.model.load_state_dict(
                maybe_wrap(
                    checkpoint["ema_state_dict" if EMA else "model_state_dict"],
                    num_processes,
                )
            )

    def eval(self):
        self.diffusion.eval()

    def train(self):
        self.diffusion.train()

    def prepare(self, objects):
        return self.accelerator.prepare(*objects)

    def train_loop(self, opt):
        print("train_dataset = FineDance_Dataset ")  
        train_dataset = FineDance_Smpl(
            args=opt,            # data/
            istrain=True,
        )
        test_dataset = FineDance_Smpl(
            args=opt,          
            istrain=False,
        )
      
        num_cpus = multiprocessing.cpu_count()
        print("batchsize=:", opt.batch_size)
        train_data_loader = DataLoader(
            train_dataset,
            batch_size=opt.batch_size,
            shuffle=True,
            num_workers=min(int(num_cpus * 0.5), 40),      # num_workers=min(int(num_cpus * 0.75), 32), 
            pin_memory=True,
            drop_last=True,
        )
        test_data_loader = DataLoader(
            test_dataset,
            batch_size=opt.batch_size,
            shuffle=True,
            num_workers=2,
            pin_memory=True,
            drop_last=True,
        )

        train_data_loader = self.accelerator.prepare(train_data_loader)
        # boot up multi-gpu training. test dataloader is only on main process
        load_loop = (
            partial(tqdm, position=1, desc="Batch")
            if self.accelerator.is_main_process
            else lambda x: x
        )
        if self.accelerator.is_main_process:
            save_dir = str(increment_path(Path(opt.project) / opt.exp_name))
            opt.exp_name = save_dir.split("/")[-1]
            wandb.init(project=opt.wandb_pj_name, name=opt.exp_name)
            save_dir = Path(save_dir)
            wdir = save_dir / "weights"
            wdir.mkdir(parents=True, exist_ok=True)
            wandb.save("params.yaml")  # 保存wandb配置到文件
            yaml_path = os.path.join(wdir, 'parameters.yaml')
            save_arguments_to_yaml(opt, yaml_path)


        self.accelerator.wait_for_everyone()
        for epoch in range(1, opt.epochs + 1):
            print("epoch:", epoch+self.resume_num)
            avg_loss = 0
            avg_vloss = 0
            avg_fkloss = 0
            avg_footloss = 0

            # train
            self.train()
            for step, (x, cond, filename) in enumerate(
                load_loop(train_data_loader)
            ):
                if opt.nfeats == 139 or opt.nfeats==135:
                    x = x[:, :, :139]

                total_loss, (loss, v_loss, fk_loss, foot_loss) = self.diffusion(
                    x, cond, t_override=None
                )
                # print("3")
                self.optim.zero_grad()
                self.accelerator.backward(total_loss)
                self.optim.step()

                # ema update and train loss update only on main
                if self.accelerator.is_main_process:
                    avg_loss += loss.detach().cpu().numpy()
                    avg_vloss += v_loss.detach().cpu().numpy()
                    avg_fkloss += fk_loss.detach().cpu().numpy()
                    avg_footloss += foot_loss.detach().cpu().numpy()
                    if step % opt.ema_interval == 0:
                        self.diffusion.ema.update_model_average(
                            self.diffusion.master_model, self.diffusion.model
                        )
                        
            #-----------------------------------------------------------------------------------------------------------
            # test
            # Save model
        
            if ((epoch+self.resume_num) % opt.save_interval) == 0  or epoch<=1:
                # everyone waits here for the val loop to finish ( don't start next train epoch early)
                self.accelerator.wait_for_everyone()
                self.eval()     # debug!
                # save only if on main thread
                if self.accelerator.is_main_process:
                    # self.eval()
                    # log
                    avg_loss /= len(train_data_loader)
                    avg_vloss /= len(train_data_loader)
                    avg_fkloss /= len(train_data_loader)
                    avg_footloss /= len(train_data_loader)
                    log_dict = {
                        "Train Loss": avg_loss,
                        "V Loss": avg_vloss,
                        "FK Loss": avg_fkloss,
                        "Foot Loss": avg_footloss,
                    }
            
                    wandb.log(log_dict)
                
                    ckpt = {
                        "ema_state_dict": self.diffusion.master_model.state_dict(),     # 经过accelerate prepare的模型，在保存时需要unwrap，反之不需要
                        "model_state_dict": self.accelerator.unwrap_model(
                            self.model
                        ).state_dict(),
                        "optimizer_state_dict": self.optim.state_dict(),
                        "normalizer": self.normalizer,
                    }
                    
                    torch.save(ckpt, os.path.join(wdir, f"train-{epoch+self.resume_num}.pt"))
                    print(f"[MODEL SAVED at Epoch {epoch+self.resume_num}]")
                   
                    # generate a sample
                    render_count = 2
                    shape = (render_count, self.horizon, self.opt.nfeats)
                    print("Generating Sample")
                    # draw a music from the test dataset
                    (x, cond, filename) = next(iter(test_data_loader))
                    # if opt.do_normalize:
                    #     x = self.normalizer.normalize(x)
                            
                    if opt.nfeats == 139 or opt.nfeats==135:
                        x = x[:, :, :139]
                        
                    cond = cond.to(self.accelerator.device)
                    # name_iter = name_iter+1
                    self.diffusion.render_sample(
                        shape,
                        cond[:render_count],
                        self.normalizer,
                        epoch+self.resume_num,
                        render_out = os.path.join(opt.render_dir, "train_" + opt.exp_name),      # render out
                        fk_out = os.path.join(opt.render_dir, "train_" + opt.exp_name),
                        name=filename[:render_count],
                        # name = str(epoch) + str(name_iter).zfill(3)
                        sound=True,
                    )
            #-----------------------------------------------------------------------------------------------------------
            
            
        if self.accelerator.is_main_process:
            wandb.run.finish()

    def render_sample(
        self, data_tuple, label, render_dir, render_count=-1, mode='normal', fk_out=None, render=True,
    ):
        _, cond, wavname = data_tuple
        assert len(cond.shape) == 3
        if render_count < 0:
            render_count = len(cond)
        shape = (render_count, self.horizon, self.repr_dim)
        cond = cond.to(self.accelerator.device).float()
        self.diffusion.render_sample(
            shape,
            cond[:render_count],
            self.normalizer,
            label,
            render_dir,
            name=wavname[:render_count],
            sound=True,
            mode=mode,           
            fk_out=fk_out,
            render=render
        )

def train(opt):
    model = EDGE(opt, opt.feature_type)
    model.train_loop(opt)
    
if __name__ == "__main__":
    opt = FineDance_parse_train_opt()
    command = ' '.join(sys.argv)
    if not os.path.exists(os.path.join(opt.project, opt.exp_name)):
        os.makedirs(os.path.join(opt.project, opt.exp_name), exist_ok=False)
    with open(os.path.join(opt.project, opt.exp_name, 'command.txt'), 'w') as f:
        f.write(command)

    yaml_path = os.path.join(opt.project, opt.exp_name, 'parameters.yaml')
    save_arguments_to_yaml(opt, yaml_path)
            
    train(opt)