train_mutiGPUs.py

import time
import argparse
import json
import logging
import math
import os
from tqdm import tqdm
from torch.utils.tensorboard import SummaryWriter
import datasets
import numpy as np
import pandas as pd
# import wandb
import torch
from torch.utils.data.distributed import DistributedSampler
from datasets import load_dataset
from torch.utils.data import Dataset, DataLoader
from tqdm.auto import tqdm
import torch.multiprocessing as mp
from tools.torch_tools import get_encode_text, get_latent
import diffusers
import transformers
from models import build_pretrained_models, AudioGPTDiffusion
from transformers import SchedulerType, get_scheduler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group, get_backend, barrier

class Trainer:
    def __init__(
        self,
        model: torch.nn.Module,
        train_data: DataLoader,
        optimizer: torch.optim.Optimizer,
        lr_scheduler: torch.optim.lr_scheduler.LambdaLR,
        gpu_id: int,
        save_every: int,
    ) -> None:
        self.gpu_id = gpu_id
        self.model = model.to(gpu_id)
        self.train_data = train_data
        self.optimizer = optimizer
        self.save_every = save_every
        self.lr_scheduler = lr_scheduler
        self.model = DDP(model, device_ids=[gpu_id])
        self.mae_path = '/home/huangqiaochu/dtj/data/audiocaps/crossattn_audiomae_generated/'
        # self.t5_path = '/home/huangqiaochu/dtj/data/audiocaps/crossattn_flan_t5/'
        self.latent_path = '/home/huangqiaochu/dtj/data/audiocaps/latents/'
        
    def _run_batch(self, text_emd, mask, latent):
        loss = self.model(latent, text_emd, mask, validation_mode=False)
        self.total_loss += loss.detach().float()
        loss.backward()
        self.optimizer.step()
        self.lr_scheduler.step()
        self.optimizer.zero_grad()

    def _run_epoch(self, epoch):
        b_sz = len(next(iter(self.train_data))[0])
        self.total_loss = 0
        print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)*epoch}")
        self.train_data.sampler.set_epoch(epoch)
        progress_bar = tqdm(range(len(self.train_data)), disable= self.gpu_id != 0)
        for text, audios, _ in self.train_data:
            for i, audio in enumerate(audios):
                name ='Y'+ '_'.join(audio.split('/')[-1].split('_')[:-1]) + '.wav'
                text[i] = name.split('/')[-1].split('.')[0]   
            text_emd, mask = get_encode_text(text, self.mae_path+'train')
            latent = get_latent(text, self.latent_path+'train')
                
            text_emd = text_emd.to(self.gpu_id)
            mask = mask.to(self.gpu_id)
            latent = latent.to(self.gpu_id)
            self._run_batch(text_emd, mask, latent)
            progress_bar.update(1)

    def _save_checkpoint(self, epoch):
        ckp = self.model.module.state_dict()
        PATH = "checkpoint.pt"
        torch.save(ckp, PATH)
        print(f"Epoch {epoch} | Training checkpoint saved at {PATH}")

    def train(self, max_epochs: int):
        for epoch in range(max_epochs):
            self._run_epoch(epoch)
            if self.gpu_id == 0 and epoch % self.save_every == 0:
                self._save_checkpoint(epoch)
                result = {}
                result["epoch"] = epoch,
                result["step"] = epoch*len(self.train_data)
                result["train_loss"] = round(self.total_loss.item()/len(self.train_data), 4)
                logging.info(result)
                with open("{}/summary.jsonl".format(args.output_dir), "a") as f:
                    f.write(json.dumps(result) + "\n\n")
                    
def parse_args(parser):
    parser.add_argument(
        "--train_file", type=str, default="data/train_audiocaps.json",
        help="A csv or a json file containing the training data."
    )
    parser.add_argument(
        "--validation_file", type=str, default="data/valid_audiocaps.json",
        help="A csv or a json file containing the validation data."
    )
    parser.add_argument(
        "--test_file", type=str, default="data/test_audiocaps.json",
        help="A csv or a json file containing the test data for generation."
    )
    parser.add_argument(
        "--num_examples", type=int, default=-1,
        help="How many examples to use for training and validation.",
    )
    
    parser.add_argument(
        "--scheduler_name", type=str, default="/home/huangqiaochu/dtj/tango/huggingface/stabilityai--stable-diffusion-2-1/scheduler_config.json",
        help="Scheduler identifier.",
    )
    parser.add_argument(
        "--unet_model_name", type=str, default=None,
        help="UNet model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--unet_model_config", type=str, default='configs/diffusion_model_config_MAE.json',
        help="UNet model config json path.",
    )
    parser.add_argument(
        "--hf_model", type=str, default=None,
        help="Tango model identifier from huggingface: declare-lab/tango",
    )
    parser.add_argument(
        "--snr_gamma", type=float, default=None,
        help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. "
        "More details here: https://arxiv.org/abs/2303.09556.",
    )
    parser.add_argument(
        "--freeze_text_encoder", action="store_true", default=True,
        help="Freeze the text encoder model.",
    )
    parser.add_argument(
        "--text_column", type=str, default="captions",
        help="The name of the column in the datasets containing the input texts.",
    )
    parser.add_argument(
        "--audio_column", type=str, default="location",
        help="The name of the column in the datasets containing the audio paths.",
    )
    parser.add_argument(
        "--augment", action="store_true", default=False,
        help="Augment training data.",
    )
    parser.add_argument(
        "--uncondition", action="store_true", default=False,
        help="10% uncondition for training.",
    )
    parser.add_argument(
        "--prefix", type=str, default=None,
        help="Add prefix in text prompts.",
    )
    parser.add_argument(
        "--per_device_train_batch_size", type=int, default=2,
        help="Batch size (per device) for the training dataloader.",
    )
    parser.add_argument(
        "--per_device_eval_batch_size", type=int, default=2,
        help="Batch size (per device) for the validation dataloader.",
    )
    parser.add_argument(
        "--learning_rate", type=float, default=3e-5,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument(
        "--weight_decay", type=float, default=1e-8,
        help="Weight decay to use."
    )
    parser.add_argument(
        "--num_train_epochs", type=int, default=40,
        help="Total number of training epochs to perform."
    )
    parser.add_argument(
        "--max_train_steps", type=int, default=None,
        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
    )
    parser.add_argument(
        "--gradient_accumulation_steps", type=int, default=4,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--lr_scheduler_type", type=SchedulerType, default="linear",
        help="The scheduler type to use.",
        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
    )
    parser.add_argument(
        "--num_warmup_steps", type=int, default=0,
        help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument(
        "--adam_beta1", type=float, default=0.9,
        help="The beta1 parameter for the Adam optimizer."
    )
    parser.add_argument(
        "--adam_beta2", type=float, default=0.999,
        help="The beta2 parameter for the Adam optimizer."
    )
    parser.add_argument(
        "--adam_weight_decay", type=float, default=1e-2,
        help="Weight decay to use."
    )
    parser.add_argument(
        "--adam_epsilon", type=float, default=1e-08,
        help="Epsilon value for the Adam optimizer"
    )
    parser.add_argument(
        "--output_dir", type=str, default=None,
        help="Where to store the final model."
    )
    parser.add_argument(
        "--checkpointing_steps", type=str, default="best",
        help="Whether the various states should be saved at the end of every 'epoch' or 'best' whenever validation loss decreases.",
    )
    parser.add_argument(
        "--resume_from_checkpoint", type=str, default=None,
        help="If the training should continue from a local checkpoint folder.",
    )
    parser.add_argument(
        "--with_tracking", action="store_true",
        help="Whether to enable experiment trackers for logging.",
    )
    parser.add_argument(
        "--report_to", type=str, default="all",
        help=(
            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
            ' `"wandb"`, `"comet_ml"` and `"clearml"`. Use `"all"` (default) to report to all integrations.'
            "Only applicable when `--with_tracking` is passed."
        ),
    )

def get_logger(__name__):
    # 创建一个日志记录器对象
    logger = logging.getLogger(__name__)

    # 设置日志级别（可根据需要设置不同的级别）
    logger.setLevel(logging.INFO)

    # 创建一个文件处理器，用于将日志信息存储到文件
    file_handler = logging.FileHandler('logs/ddp.log')
    file_handler.setLevel(logging.INFO)

    # 创建一个控制台处理器，用于将日志信息打印到控制台
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    # 创建一个日志格式化器，用于指定日志信息的格式
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)

    # 将处理器添加到日志记录器对象
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)

    return logger

def ddp_setup(rank, world_size):
    """
    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    
    init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
   
def save_args(args):
    # Handle output directory creation and wandb tracking, 只在主进程设置
    if args.output_dir is None or args.output_dir == "":
        args.output_dir = "saved/" + str(int(time.time()))

    os.makedirs("{}/{}".format(args.output_dir, "outputs"), exist_ok=True)
    with open("{}/summary.jsonl".format(args.output_dir), "a") as f:
        f.write(json.dumps(dict(vars(args))) + "\n\n")

class Text2AudioDataset(Dataset):
    def __init__(self, dataset, prefix, text_column, audio_column, num_examples=-1):

        inputs = list(dataset[text_column])
        self.inputs = [prefix + inp for inp in inputs]
        self.audios = list(dataset[audio_column])
        self.indices = list(range(len(self.inputs)))

        self.mapper = {}
        for index, audio, text in zip(self.indices, self.audios, inputs):
            self.mapper[index] = [audio, text]

        if num_examples != -1:
            self.inputs, self.audios = self.inputs[:num_examples], self.audios[:num_examples]
            self.indices = self.indices[:num_examples]

    def __len__(self):
        return len(self.inputs)

    def get_num_instances(self):
        return len(self.inputs)

    def __getitem__(self, index):
        s1, s2, s3 = self.inputs[index], self.audios[index], self.indices[index]
        return s1, s2, s3

    def collate_fn(self, data):
        dat = pd.DataFrame(data)
        return [dat[i].tolist() for i in dat]

def prepare_dataset(args):
    # Get the datasets
    data_files = {}
    if args.train_file is not None:
        data_files["train"] = args.train_file
    if args.validation_file is not None:
        data_files["validation"] = args.validation_file

    if args.test_file is not None:
        data_files["test"] = args.test_file
    else:
        if args.validation_file is not None:
            data_files["test"] = args.validation_file

    extension = args.train_file.split(".")[-1]
    raw_datasets = load_dataset(extension, data_files=data_files)
    
    if args.prefix:
        prefix = args.prefix
    else:
        prefix = ""
    text_column, audio_column = args.text_column, args.audio_column
    train_dataset = Text2AudioDataset(raw_datasets["train"], prefix, text_column, audio_column, args.num_examples)
    eval_dataset = Text2AudioDataset(raw_datasets["validation"], prefix, text_column, audio_column, args.num_examples)
    test_dataset = Text2AudioDataset(raw_datasets["test"], prefix, text_column, audio_column, args.num_examples)

    
    return train_dataset, eval_dataset, test_dataset
  
def load_train_objs(rank, args):
    
 
    
    if rank == 0:
        train_dataset, eval_dataset, test_dataset = prepare_dataset(args)
        print("Num instances in train: {}, validation: {}, test: {}".format(train_dataset.get_num_instances(), eval_dataset.get_num_instances(), test_dataset.get_num_instances()))
    barrier()
    
    if rank != 0:
        train_dataset, eval_dataset, test_dataset = prepare_dataset(args)
    barrier()
    
    model = AudioGPTDiffusion(
        args.scheduler_name, args.unet_model_name, args.unet_model_config, args.snr_gamma, args.freeze_text_encoder, args.uncondition
    )
    
    # Optimizer
    if args.unet_model_config:
        optimizer_parameters = model.unet.parameters()
        
    num_trainable_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    if rank == 0:
        print("Optimizing UNet parameters.")
        print("Num trainable parameters: {}".format(num_trainable_parameters))

    optimizer = torch.optim.AdamW(
        optimizer_parameters, lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )
    


    return train_dataset, eval_dataset, test_dataset, model, optimizer

def prepare_dataloader(dataset: Dataset, batch_size: int):
    return DataLoader(
        dataset,
        batch_size=batch_size,
        pin_memory=True,
        shuffle=False,
        sampler=DistributedSampler(dataset)
    )
  
def main(rank: int, world_size: int, args: argparse.Namespace):
    
    torch.manual_seed(args.seed)
    ddp_setup(rank, world_size)
     # Log the process state
    logger = get_logger("rank_"+str(rank))
    logger.info(f"Process rank: {rank}, World size: {world_size}, Process group: {get_backend()}")

        
    train_dataset, eval_dataset, test_dataset, model, optimizer = load_train_objs(rank, args)
    train_dataloader = prepare_dataloader(train_dataset, args.batch_size)
    
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) )
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps, 
        num_training_steps=args.max_train_steps 
    )
        
    trainer = Trainer(model, train_dataloader, optimizer, lr_scheduler, rank, args.save_every)
    ### Train!
    if rank == 0:
        save_args(args)
        total_batch_size = args.batch_size * world_size 
        logger.info("***** Running training *****")
        logger.info(f"  Num Epochs = {args.total_epochs}")
        logger.info(f"  Instantaneous batch size per device = {args.batch_size}")
        logger.info(f"  Total train batch size (w. parallel, distributed) = {total_batch_size}")
        
    trainer.train(args.total_epochs)
    destroy_process_group()
    
    
            
if __name__ == "__main__":
    
    parser = argparse.ArgumentParser(description='simple distributed training job')
    parser.add_argument('--total_epochs', default=40, type=int, help='Total epochs to train the model')
    parser.add_argument('--save_every', default=5, type=int, help='How often to save a snapshot')
    parser.add_argument('--batch_size', default=1, type=int, help='Input batch size on each device (default: 32)')
    parser.add_argument('--seed', default=42, type=int, help='Random seed (default: 42)')
    parse_args(parser) ## add more args
    args = parser.parse_args()
    
    
    
    world_size = torch.cuda.device_count()
    mp.spawn(main, args=(world_size, args), nprocs=world_size)