train_hf.py

import argparse
import functools
import logging
import math
import os
from datetime import timedelta

import datasets
import torch
import torch.nn.functional as F
import torch.distributed as dist
from accelerate import Accelerator, InitProcessGroupKwargs
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed

from datasets import load_dataset
from torchvision import transforms
from torchvision.datasets import ImageFolder

import diffusers
from diffusers.training_utils import compute_snr
from diffusers.optimization import get_scheduler

from models import All_models, DiT, Transformer, EMAModel
from timm.models import create_model
from utils import center_crop_arr, safe_blob_write, load_vae
from schedule.ddpm import DDPMScheduler

import wandb

logger = get_logger(__name__, log_level="INFO")


def parse_args():  
    parser = argparse.ArgumentParser(description="Simple example of a training script.")  
  
    # 基本参数  
    parser.add_argument("--seed", type=int, default=0, help="A seed to use for the random number generator. Can be negative to not set a seed.")  
    parser.add_argument("--output_dir", type=str, default="results", help="The output directory where the model predictions and checkpoints will be written.")  
    parser.add_argument("--cache_dir", type=str, default="/mnt/msranlp/yutao/cache", help="The directory where the downloaded models and datasets will be stored.")  
    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")  
  
    # 数据集参数  
    parser.add_argument("--dataset_name", type=str, default=None, help="The name of the Dataset (from the HuggingFace hub) to train on.")  
    parser.add_argument("--dataset_config_name", type=str, default=None, help="The config of the Dataset, leave as None if there's only one config.")  
    parser.add_argument("--train_data_dir", type=str, default="/tmp/ILSVRC/Data/CLS-LOC/train", help="A folder containing the training data.")  
      
    # 模型参数  
    parser.add_argument("--model", type=str, default="Transformer-L", help="The config of the UNet model to train.")  
    parser.add_argument("--vae", type=str, default=None, help="Path to pre-trained VAE model.")  
    parser.add_argument("--image_size", type=int, default=256, help="The image_size for input images.")  
    parser.add_argument("--num_classes", type=int, default=1000, help="Number of classes for the model.")
    parser.add_argument("--dropout", type=float, default=0.0, help="Dropout probability.")  
  
    # 训练参数  
    parser.add_argument("--batch_size", type=int, default=32, help="Batch size (per device) for the training dataloader.")  
    parser.add_argument("--num_epochs", type=int, default=100, help="Number of epochs to train for.")  
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.")  
    parser.add_argument("--dataloader_num_workers", type=int, default=2, help="The number of subprocesses to use for data loading.")  
      
    # 优化器参数  
    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.")  
    parser.add_argument("--lr_scheduler", type=str, default="cosine", help="The scheduler type to use.")  
    parser.add_argument("--lr_warmup_steps", type=int, default=100, 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.98, help="The beta2 parameter for the Adam optimizer.")  
    parser.add_argument("--adam_weight_decay", type=float, default=0.01, help="Weight decay magnitude for the Adam optimizer.")  
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer.")  
      
    # EMA参数  
    parser.add_argument("--use_ema", action="store_true", help="Whether to use Exponential Moving Average for the final model weights.")  
    parser.add_argument("--ema_inv_gamma", type=float, default=1.0, help="The inverse gamma value for the EMA decay.")  
    parser.add_argument("--ema_power", type=float, default=3 / 4, help="The power value for the EMA decay.")  
    parser.add_argument("--ema_max_decay", type=float, default=0.9999, help="The maximum decay magnitude for EMA.")  
      
    # 日志参数  
    parser.add_argument("--logger", type=str, default=None, help="The logger type to use.")  
    parser.add_argument("--logging_dir", type=str, default="logs", help="The directory to store logs.")  
    parser.add_argument("--wandb_project", type=str, default=None, help="The wandb project name.")  
    parser.add_argument("--wandb_entity", type=str, default=None, help="The wandb entity (username or team).")  
    parser.add_argument("--log_every", type=int, default=100, help="Log every X steps.")  
      
    # 分布式训练参数  
    parser.add_argument("--mixed_precision", type=str, default="no", choices=["no", "fp16", "bf16"], help="Whether to use mixed precision.")  

    # DDPM参数  
    parser.add_argument("--prediction_type", type=str, default="epsilon", help="Whether the model should predict the 'epsilon'/noise error or directly the reconstructed image 'x0'.")  
    parser.add_argument("--ddpm_num_steps", type=int, default=1000, help="The number of steps to use for DDPM.")  
    parser.add_argument("--ddpm_num_inference_steps", type=int, default=1000, help="The number of inference steps to use for DDPM.")
    parser.add_argument("--ddpm_beta_schedule", type=str, default="cosine", help="The beta schedule to use for DDPM.") 
    parser.add_argument("--ddpm_batch_mul", type=int, default=4, help="The batch multiplier to use for DDPM.")  
    parser.add_argument("--checkpointing_steps", type=int, default=5000, help="Save a checkpoint of the training state every X updates.")  
    parser.add_argument("--checkpoint", type=str, default=None, help="Resume training from a previous checkpoint.")  
      
    args = parser.parse_args()
      
    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))  
    if env_local_rank != -1 and env_local_rank != args.local_rank:  
        args.local_rank = env_local_rank  
    if args.dataset_name is None and args.train_data_dir is None:  
        raise ValueError("You must specify either a dataset name from the hub or a train data directory.")  
      
    return args  


def main(args):
    set_seed(args.seed)
    logging_dir = os.path.join(args.output_dir, args.logging_dir)
    
    vae, input_size, latent_size, flatten_input = load_vae(args.vae, args.image_size)
        
    model = All_models[args.model](
        input_size=input_size,
        in_channels=latent_size,
        num_classes=args.num_classes,
        flatten_input=flatten_input,
        drop=args.dropout,
    )
    if args.mixed_precision == "bf16":
        dtype = torch.bfloat16
    elif args.mixed_precision == "fp16":
        dtype = torch.float16
    else:
        dtype = torch.float32

    # Create EMA for the model.
    if args.use_ema:
        ema_model = EMAModel(
            model.parameters(),
            decay=args.ema_max_decay,
            min_decay=args.ema_max_decay,
            use_ema_warmup=True,
            inv_gamma=args.ema_inv_gamma,
            power=args.ema_power,
        )
    # Initialize the scheduler
    noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule, prediction_type=args.prediction_type)

    # Initialize the optimizer
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    # Initialize the accelerator
    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
    kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=7200))  # a big number for high image_size or big dataset
    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
        log_with=args.logger,
        project_config=accelerator_project_config,
        kwargs_handlers=[kwargs],
    )
    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        datasets.utils.logging.set_verbosity_warning()
        diffusers.utils.logging.set_verbosity_info()
    else:
        datasets.utils.logging.set_verbosity_error()
        diffusers.utils.logging.set_verbosity_error()

    logger.info(args)
    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
        if args.wandb_project is not None:
            wandb.init(project=args.wandb_project, entity=args.wandb_entity, config=args)  

    logger.info(model)
    logger.info(f"Model Parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    # In distributed training, the load_dataset function guarantees that only one local process can concurrently
    # download the dataset.
    augmentations = transforms.Compose([
        transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.image_size)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
    ])
    if args.dataset_name is not None:
        dataset = load_dataset(
            args.dataset_name,
            args.dataset_config_name,
            cache_dir=args.cache_dir,
            split="train",
        )
        def transform_images(examples):
            images = [augmentations(image.convert("RGB")) for image in examples["image"]]
            return {"input": images}
        dataset.set_transform(transform_images)
    else:
        dataset = ImageFolder(args.train_data_dir, transform=augmentations)
    
    train_dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.dataloader_num_workers
    )
    # Initialize the learning rate scheduler
    lr_scheduler = get_scheduler(
        args.lr_scheduler,
        optimizer=optimizer,
        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
        num_training_steps=(len(train_dataloader) * args.num_epochs // args.gradient_accumulation_steps),
    )
    # Prepare everything with our `accelerator`.
    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, lr_scheduler
    )
    # vae = accelerator.prepare_model(vae, evaluation_mode=True, device_placement=True)
    vae.to(accelerator.device)
    vae.eval()
    if args.use_ema:
        ema_model.to(accelerator.device)

    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
    if accelerator.is_main_process:
        run = os.path.split(__file__)[-1].split(".")[0]
        accelerator.init_trackers(run)

    total_batch_size = args.batch_size * accelerator.num_processes * args.gradient_accumulation_steps
    max_train_steps = len(train_dataloader) * args.num_epochs // args.gradient_accumulation_steps

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(dataset)}")
    logger.info(f"  Num Epochs = {args.num_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    global_step = 0
    running_loss = 0
    first_epoch = 0
    scaling_factor = None
    bias_factor = None
    # Potentially load in the weights and states from a previous save
    checkpoint_path = args.checkpoint
    if checkpoint_path is None and os.path.exists(os.path.join(args.output_dir, "latest")):
        with open(os.path.join(args.output_dir, "latest"), "r") as f:
            checkpoint_path = f.read().strip()
        
    if checkpoint_path is not None:
        accelerator.print(f"Resuming from checkpoint {checkpoint_path}")
        accelerator.load_state(checkpoint_path)
        other_state = torch.load(os.path.join(checkpoint_path, "other_state.pth"), map_location="cpu")
        global_step = other_state["steps"]
        scaling_factor = other_state["scaling_factor"]
        bias_factor = other_state["bias_factor"]
        if args.use_ema:
            ema_model.load_state_dict(other_state["ema"])
            logger.info("EMA model loaded successfully")
        first_epoch = global_step * args.gradient_accumulation_steps // len(train_dataloader)
        resume_step = global_step * args.gradient_accumulation_steps % len(train_dataloader)

    # Train!
    # snr = compute_snr(noise_scheduler, torch.arange(args.ddpm_num_steps, device=accelerator.device))
    # sample_weight = (
    #     torch.stack([snr, 5 * torch.ones(args.ddpm_num_steps, device=accelerator.device)], dim=1).min(dim=1)[0] / snr
    # )
    sample_weight = torch.ones(args.ddpm_num_steps, device=accelerator.device)
    for epoch in range(first_epoch, args.num_epochs):
        model.train()
        for step, (clean_images, label) in enumerate(train_dataloader):
            # Skip steps until we reach the resumed step
            if args.checkpoint and epoch == first_epoch:
                if step < resume_step:
                    continue
                
            with torch.no_grad():
                vae_latent = vae.encode(clean_images)
                clean_images = vae_latent.sample()
                mode_images = vae_latent.mode()
                if scaling_factor is None:
                    scaling_factor = 1. / clean_images.flatten().std()
                    bias_factor = -clean_images.flatten().mean()
                    dist.all_reduce(scaling_factor, op=dist.ReduceOp.SUM)
                    dist.all_reduce(bias_factor, op=dist.ReduceOp.SUM)
                    scaling_factor = scaling_factor.item() / dist.get_world_size()
                    bias_factor = bias_factor.item() / dist.get_world_size()
                    logger.info(f"Scaling factor: {scaling_factor}, Bias factor: {bias_factor}")
                clean_images = (clean_images + bias_factor) * scaling_factor
                mode_images = (mode_images + bias_factor) * scaling_factor

            with accelerator.accumulate(model):
                bsz, latent_size, h, w = clean_images.shape
                if isinstance(model.module, Transformer):
                    noise = torch.randn((bsz * args.ddpm_batch_mul * h * w, latent_size), device=clean_images.device, dtype=clean_images.dtype)
                    timesteps = torch.multinomial(sample_weight, bsz * args.ddpm_batch_mul * h * w, replacement=True)
                    clean_images_repeated = clean_images.repeat_interleave(args.ddpm_batch_mul, dim=0).permute(0, 2, 3, 1).reshape(-1, clean_images.shape[1])
                    noisy_images = noise_scheduler.add_noise(clean_images_repeated, noise, timesteps)
                    velocity = noise_scheduler.get_velocity(clean_images_repeated, noise, timesteps)
                    noisy_images, noise, velocity = [x.reshape(bsz * args.ddpm_batch_mul, h, w, latent_size).permute(0, 3, 1, 2) for x in [noisy_images, noise, velocity]]
                    timesteps = timesteps.reshape(bsz * args.ddpm_batch_mul, h * w)
                    model_output = model(noisy_images.to(dtype), timesteps.to(dtype), x_start=clean_images.to(dtype), y=label, batch_mul=args.ddpm_batch_mul)
                elif isinstance(model.module, DiT):
                    noise = torch.randn_like(clean_images)
                    timesteps = torch.multinomial(sample_weight, bsz, replacement=True)
                    noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
                    velocity = noise_scheduler.get_velocity(clean_images, noise, timesteps)
                    model_output = model(noisy_images.to(dtype), timesteps.to(dtype), y=label)
                else:
                    raise NotImplementedError()
                if args.prediction_type == "epsilon":
                    loss = F.mse_loss(model_output.float(), noise.float())
                elif args.prediction_type == "v_prediction":
                    loss = F.mse_loss(model_output.float(), velocity.float())
                else:
                    raise NotImplementedError()
                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    gnorm = accelerator.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()                 

            running_loss += loss.item()
            if accelerator.sync_gradients:
                global_step += 1
                if args.use_ema:
                    ema_model.step(model.parameters())
                if global_step % args.log_every == 0:
                    avg_loss = running_loss / args.log_every / args.gradient_accumulation_steps
                    running_loss = 0
                    logs = {"loss": avg_loss, "lr": lr_scheduler.get_last_lr()[0], "step": global_step, "gnorm": gnorm.item(), "batch size": total_batch_size, "epoch": epoch}
                    if args.use_ema:
                        logs["ema_decay"] = ema_model.cur_decay_value
                    logger.info(logs)
                    accelerator.log(logs, step=global_step)
                    if accelerator.is_main_process and args.wandb_project is not None:  
                        wandb.log(logs, step=global_step)
                
                if global_step % args.checkpointing_steps == 0:
                    def save_checkpoint(path):
                        accelerator.save_state(path)
                        if accelerator.is_main_process:
                            other_state = {
                                "scaling_factor": scaling_factor,
                                "bias_factor": bias_factor,
                                "steps": global_step,
                                "ema": ema_model.state_dict() if args.use_ema else None,
                            }
                            torch.save(other_state, os.path.join(path, "other_state.pth"))
                    save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                    save_checkpoint(os.path.join(save_path))
                    if accelerator.is_main_process:
                        safe_blob_write(os.path.join(args.output_dir, "latest"), save_path)
                    logger.info(f"Saved state to {save_path}")

    accelerator.end_training()


if __name__ == "__main__":
    args = parse_args()
    main(args)