c2i_train_infer.py

import argparse
import datetime
import numpy as np
import os
import time
import functools
from pathlib import Path

import torch
import torch.backends.cudnn as cudnn
from torch.utils.tensorboard import SummaryWriter
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    BackwardPrefetch,
    ShardingStrategy,
    FullStateDictConfig,
    StateDictType,
)
from torch.distributed.fsdp.wrap import (
    transformer_auto_wrap_policy,
    enable_wrap,
    wrap,
)

from grn.utils_c2i.crop import center_crop_arr
import grn.utils_c2i.misc as misc

import copy
from grn.utils_c2i.engine import train_one_epoch, evaluate
from grn.utils import wandb_utils as wandb_utils

from grn.utils_c2i.denoiser import Denoiser
from grn.models.grn_c2i import GRNblock


def get_args_parser():
    parser = argparse.ArgumentParser('GRN', add_help=False)

    # architecture
    parser.add_argument('--model', default='GRN_B', type=str, metavar='MODEL',
                        help='Name of the model to train')
    parser.add_argument('--img_size', default=256, type=int, help='Image size')
    parser.add_argument('--attn_dropout', type=float, default=0.0, help='Attention dropout rate')
    parser.add_argument('--proj_dropout', type=float, default=0.0, help='Projection dropout rate')

    # training
    parser.add_argument('--epochs', default=200, type=int)
    parser.add_argument('--warmup_epochs', type=int, default=5, metavar='N',
                        help='Epochs to warm up LR')
    parser.add_argument('--batch_size', default=128, type=int,
                        help='Batch size per GPU (effective batch size = batch_size * # GPUs)')
    parser.add_argument('--lr', type=float, default=None, metavar='LR',
                        help='Learning rate (absolute)')
    parser.add_argument('--blr', type=float, default=5e-5, metavar='LR',
                        help='Base learning rate: absolute_lr = base_lr * total_batch_size / 256')
    parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                        help='Minimum LR for cyclic schedulers that hit 0')
    parser.add_argument('--lr_schedule', type=str, default='constant',
                        help='Learning rate schedule')
    parser.add_argument('--weight_decay', type=float, default=0.0,
                        help='Weight decay (default: 0.0)')
    parser.add_argument('--ema_decay1', type=float, default=0.9999,
                        help='The first ema to track. Use the first ema for sampling by default.')
    parser.add_argument('--ema_decay2', type=float, default=0.9996,
                        help='The second ema to track')
    parser.add_argument('--P_mean', default=-0.8, type=float)
    parser.add_argument('--P_std', default=0.8, type=float)
    parser.add_argument('--noise_scale', default=1.0, type=float)
    parser.add_argument('--t_eps', default=5e-2, type=float)
    parser.add_argument('--label_drop_prob', default=0.1, type=float)

    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                        help='Starting epoch')
    parser.add_argument('--num_workers', default=12, type=int)
    parser.add_argument('--pin_mem', action='store_true',
                        help='Pin CPU memory in DataLoader for faster GPU transfers')
    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
    parser.set_defaults(pin_mem=True)

    # sampling
    parser.add_argument('--sampling_method', default='heun', type=str,
                        help='ODE samping method')
    parser.add_argument('--num_sampling_steps', default=50, type=int,
                        help='Sampling steps')
    parser.add_argument('--cfg', default=1.0, type=float,
                        help='Classifier-free guidance factor')
    parser.add_argument('--interval_min', default=0.0, type=float,
                        help='CFG interval min')
    parser.add_argument('--interval_max', default=1.0, type=float,
                        help='CFG interval max')
    parser.add_argument('--num_images', default=50000, type=int,
                        help='Number of images to generate')
    parser.add_argument('--eval_freq', type=int, default=40,
                        help='Frequency (in epochs) for evaluation')
    parser.add_argument('--online_eval', type=int, default=0, choices=[0,1],
                        help='Whether to evaluate the model online')
    parser.add_argument('--evaluate_gen', action='store_true')
    parser.add_argument('--gen_bsz', type=int, default=256,
                        help='Generation batch size')

    # dataset
    parser.add_argument('--data_path', default='./data/imagenet', type=str,
                        help='Path to the dataset')
    parser.add_argument('--class_num', default=1000, type=int)

    # checkpointing
    parser.add_argument('--output_dir', default='./output_dir',
                        help='Directory to save outputs (empty for no saving)')
    parser.add_argument('--resume', default='',
                        help='Folder that contains checkpoint to resume from')
    parser.add_argument('--save_last_freq', type=int, default=5,
                        help='Frequency (in epochs) to save checkpoints')
    parser.add_argument('--log_freq', default=100, type=int)
    parser.add_argument('--device', default='cuda',
                        help='Device to use for training/testing')

    # distributed training
    parser.add_argument('--world_size', default=1, type=int,
                        help='Number of distributed processes')
    parser.add_argument('--local_rank', default=-1, type=int)
    parser.add_argument('--dist_on_itp', action='store_true')
    parser.add_argument('--dist_url', default='env://',
                        help='URL used to set up distributed training')
    parser.add_argument('--hbq_round', default=4, type=int,)
    parser.add_argument('--in_channels', default=3, type=int,)
    parser.add_argument('--method', default='GRN_ind', type=str, choices=['GRN_ind', 'GRN_bit'])
    parser.add_argument('--vae_path', default='', type=str,)
    parser.add_argument('--tau', default=1.0, type=float,)
    parser.add_argument('--wandb', default=1, type=int, choices=[0,1])
    parser.add_argument('--generation_dir', default='/tmp', type=str)
    parser.add_argument('--clip_grad_norm', default=1., type=float)
    parser.add_argument('--use_fsdp_train', default=0, type=int, choices=[0, 1])
    parser.add_argument('--delete_images', default=1, type=int, choices=[0, 1])
    parser.add_argument('--use_confidence_sampling', default=0, type=int, choices=[0, 1])
    parser.add_argument('--inner_shard_degree', default=8, type=int)
    parser.add_argument('--patch_size', default=1, type=int)
    parser.add_argument('--convert_type', default='', type=str)
    parser.add_argument('--mask_group_size', default=-1, type=int)
    parser.add_argument('--grn_shift_factor', default=1., type=float)
    parser.add_argument('--use_focal_loss', default=0, type=int, choices=[0, 1])
    return parser


def main(args):
    misc.init_distributed_mode(args)
    print('Job directory:', os.path.dirname(os.path.realpath(__file__)))
    print("Arguments:\n{}".format(args).replace(', ', ',\n'))

    device = torch.device(args.device)

    # Set seeds for reproducibility
    seed = args.seed + misc.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)

    cudnn.benchmark = True

    num_tasks = misc.get_world_size()
    global_rank = misc.get_rank()

    # Set up TensorBoard logging (only on main process)
    if global_rank == 0 and args.output_dir is not None:
        os.makedirs(args.output_dir, exist_ok=True)
        log_writer = SummaryWriter(log_dir=args.output_dir)
        if args.wandb:
            entity = os.environ["EXP_NAME"]
            project = os.environ["PROJECT"]
            wandb_utils.wandb.init(project=project, name=entity, config={})
    else:
        log_writer = None

    # Data augmentation transforms
    transform_train = transforms.Compose([
        transforms.Lambda(lambda img: center_crop_arr(img, args.img_size)),
        transforms.RandomHorizontalFlip(),
        transforms.PILToTensor()
    ])

    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, 'train'), transform=transform_train)
    print(dataset_train)

    sampler_train = torch.utils.data.DistributedSampler(
        dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
    )
    print("Sampler_train =", sampler_train)

    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, sampler=sampler_train,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=True
    )

    torch._dynamo.config.cache_size_limit = 128
    torch._dynamo.config.optimize_ddp = False

    # Create denoiser
    model = Denoiser(args)

    # ininitalize vae
    from grn.models.hbq_tokenizer import HBQ_Tokenizer
    vae = HBQ_Tokenizer(args=args, latent_channels=16, encoder_out_type='feature_tanh')
    vae.eval()
    vae = vae.to('cuda')
    for param in vae.parameters():
        param.requires_grad = False
    state_dict = torch.load(args.vae_path, map_location='cuda')
    if 'ema' in state_dict:
        print(f'Load ema vae weights')
        state_dict = state_dict['ema']
    else:
        print(f'Load non ema vae weights')
        state_dict = state_dict['vae']
    print('Load vae: ', vae.load_state_dict(state_dict, assign=True))

    print("Model =", model)
    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print("Number of trainable parameters: {:.6f}M".format(n_params / 1e6))

    model.to(device)

    eff_batch_size = args.batch_size * misc.get_world_size()
    if args.lr is None:  # only base_lr (blr) is specified
        args.lr = args.blr * eff_batch_size / 256

    print("Base lr: {:.2e}".format(args.lr * 256 / eff_batch_size))
    print("Actual lr: {:.2e}".format(args.lr))
    print("Effective batch size: %d" % eff_batch_size)

    if args.use_fsdp_train:
        auto_wrap_policy = functools.partial(
            transformer_auto_wrap_policy,
            transformer_layer_cls={GRNblock},
        )
        if args.inner_shard_degree > 0:
            sharding_strategy = ShardingStrategy.HYBRID_SHARD
            world_size = misc.get_world_size()
            assert world_size % args.inner_shard_degree == 0
            assert args.inner_shard_degree > 1 and args.inner_shard_degree <= world_size
            device_mesh = init_device_mesh('cuda', (world_size // args.inner_shard_degree, args.inner_shard_degree))
        else:
            sharding_strategy = ShardingStrategy.FULL_SHARD
            device_mesh = None
        model = FSDP(
            model,
            auto_wrap_policy=auto_wrap_policy,
            mixed_precision=MixedPrecision(
                param_dtype=torch.bfloat16, 
                reduce_dtype=torch.bfloat16, 
                buffer_dtype=torch.bfloat16
            ),
            device_id=torch.cuda.current_device(),
            sharding_strategy=sharding_strategy, 
            use_orig_params=True,
            device_mesh=device_mesh,
        )
        model_without_ddp = model
    else:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
        model_without_ddp = model.module

    # Set up optimizer with weight decay adjustment for bias and norm layers
    param_groups = misc.add_weight_decay(model_without_ddp, args.weight_decay)
    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
    print(optimizer)

    # Resume from checkpoint if provided
    # checkpoint_path = os.path.join(args.resume, "checkpoint-last.pth") if args.resume else None
    checkpoint_path = args.resume if args.resume else None
    if checkpoint_path and os.path.exists(checkpoint_path):
        checkpoint = torch.load(checkpoint_path, map_location='cpu')
        model_without_ddp.load_state_dict(checkpoint['model'])

        if args.use_fsdp_train:
            # For FSDP, load EMA state dict into model temporarily to set ema_params
            model_without_ddp.load_state_dict(checkpoint['model_ema1'])
            model_without_ddp.module.ema_params1 = [p.detach().clone() for p in model_without_ddp.parameters()]
            
            model_without_ddp.load_state_dict(checkpoint['model_ema2'])
            model_without_ddp.module.ema_params2 = [p.detach().clone() for p in model_without_ddp.parameters()]
            
            # Restore model
            model_without_ddp.load_state_dict(checkpoint['model'])
        else:
            ema_state_dict1 = checkpoint['model_ema1']
            ema_state_dict2 = checkpoint['model_ema2']
            model_without_ddp.ema_params1 = [ema_state_dict1[name].cuda() for name, _ in model_without_ddp.named_parameters()]
            model_without_ddp.ema_params2 = [ema_state_dict2[name].cuda() for name, _ in model_without_ddp.named_parameters()]
        
        print("Resumed checkpoint from", args.resume)

        try:
            if 'optimizer' in checkpoint and 'epoch' in checkpoint:
                if args.use_fsdp_train:
                    opt_state = FSDP.optim_state_dict_to_load(
                        model_without_ddp, optimizer, checkpoint['optimizer']
                    )
                    optimizer.load_state_dict(opt_state)
                else:
                    optimizer.load_state_dict(checkpoint['optimizer'])
                print("Loaded optimizer & scaler state!")
        except:
            print("Failed to load optimizer & scaler state! Just load checkpoint.")
        args.start_epoch = checkpoint['epoch'] + 1
        del checkpoint
    else:
        if args.use_fsdp_train:
            model_without_ddp.module.ema_params1 = [p.detach().clone() for p in model_without_ddp.parameters()]
            model_without_ddp.module.ema_params2 = [p.detach().clone() for p in model_without_ddp.parameters()]
        else:
            model_without_ddp.ema_params1 = [p.detach().clone() for p in model_without_ddp.parameters()]
            model_without_ddp.ema_params2 = [p.detach().clone() for p in model_without_ddp.parameters()]
        print("Training from scratch")

    # Evaluate generation
    if args.evaluate_gen:
        print("Evaluating checkpoint at {} epoch".format(args.start_epoch))
        with torch.random.fork_rng():
            torch.manual_seed(seed)
            with torch.no_grad():
                evaluate(model_without_ddp, args, args.start_epoch, batch_size=args.gen_bsz, log_writer=log_writer, vae=vae)
        return

    # Training loop
    print(f"Start training for {args.epochs} epochs")
    start_time = time.time()
    for epoch in range(args.start_epoch, args.epochs):
        if args.distributed:
            data_loader_train.sampler.set_epoch(epoch)

        train_one_epoch(model, model_without_ddp, data_loader_train, optimizer, device, epoch, log_writer=log_writer, args=args, vae=vae)

        # Save checkpoint periodically
        if epoch % args.save_last_freq == 0 or epoch + 1 == args.epochs:
            if misc.is_main_process():
                from grn.utils.safe_rm import safe_remove
                safe_remove(f'{args.output_dir}/checkpoint-tmp_*.pth', args.output_dir)
            misc.save_model(
                args=args,
                model_without_ddp=model_without_ddp,
                optimizer=optimizer,
                epoch=epoch,
                epoch_name=f"tmp_{epoch}"
            )

        if epoch % 100 == 0 and epoch > 0:
            misc.save_model(
                args=args,
                model_without_ddp=model_without_ddp,
                optimizer=optimizer,
                epoch=epoch
            )

        # Perform online evaluation at specified intervals
        if args.online_eval and (epoch % args.eval_freq == 0 or epoch + 1 == args.epochs):
            torch.cuda.empty_cache()
            with torch.no_grad():
                evaluate(model_without_ddp, args, epoch, batch_size=args.gen_bsz, log_writer=log_writer, vae=vae)
            torch.cuda.empty_cache()

        if misc.is_main_process() and log_writer is not None:
            log_writer.flush()

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time:', total_time_str)


if __name__ == '__main__':
    args = get_args_parser().parse_args()
    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    main(args)