scripts/vit_triplane_train.py

"""
Train a diffusion model on images.
"""
from pdb import set_trace as st
import random
import json
import sys
import os

sys.path.append('.')

import torch.distributed as dist

import traceback

import torch as th

# if th.cuda.is_available(): # FIXME
#     from xformers.triton import FusedLayerNorm as LayerNorm

import torch.multiprocessing as mp
import numpy as np

import argparse
import dnnlib
from guided_diffusion import dist_util, logger
from guided_diffusion.script_util import (
    args_to_dict,
    add_dict_to_argparser,
)
from nsr.train_nv_util import TrainLoop3DRecNV, TrainLoop3DRec, TrainLoop3DRecNVPatch, TrainLoop3DRecNVPatchSingleForward, TrainLoop3DRecNVPatchSingleForwardMV, TrainLoop3DRecNVPatchSingleForwardMV_NoCrop, TrainLoop3DRecNVPatchSingleForwardMVAdvLoss, TrainLoop3DRecNVPatchSingleForwardMV_NoCrop_adv
from nsr.script_util import create_3DAE_model, encoder_and_nsr_defaults, loss_defaults, rendering_options_defaults, eg3d_options_default, dataset_defaults
# from datasets import g_buffer_objaverse
from nsr.losses.builder import E3DGELossClass, E3DGE_with_AdvLoss


# th.backends.cuda.matmul.allow_tf32 = True # https://huggingface.co/docs/diffusers/optimization/fp16
th.backends.cuda.matmul.allow_tf32 = True
th.backends.cudnn.allow_tf32 = True
th.backends.cudnn.enabled = True


def training_loop(args):
    # def training_loop(args):
    dist_util.setup_dist(args)
    # th.autograd.set_detect_anomaly(True) # type: ignore
    th.autograd.set_detect_anomaly(False)  # type: ignore
    # https://blog.csdn.net/qq_41682740/article/details/126304613

    SEED = args.seed

    # dist.init_process_group(backend='nccl', init_method='env://', rank=args.local_rank, world_size=th.cuda.device_count())
    logger.log(f"global_rank={args.global_rank}, local_rank={args.local_rank} init complete, seed={SEED}")
    th.cuda.set_device(args.local_rank)
    th.cuda.empty_cache()

    # * deterministic algorithms flags
    th.cuda.manual_seed_all(SEED)
    np.random.seed(SEED)
    random.seed(SEED)

    # logger.configure(dir=args.logdir, format_strs=["tensorboard", "csv"])
    logger.configure(dir=args.logdir)

    logger.log("creating encoder and NSR decoder...")
    # device = dist_util.dev()
    device = th.device("cuda", args.local_rank)

    # shared eg3d opts
    opts = eg3d_options_default()

    if args.sr_training:
        args.sr_kwargs = dnnlib.EasyDict(
            channel_base=opts.cbase,
            channel_max=opts.cmax,
            fused_modconv_default='inference_only',
            use_noise=True
        )  # ! close noise injection? since noise_mode='none' in eg3d

    auto_encoder = create_3DAE_model(
        **args_to_dict(args,
                       encoder_and_nsr_defaults().keys()))
    auto_encoder.to(device)
    auto_encoder.train()

    logger.log("creating data loader...")
    # data = load_data(
    # st()
    if args.objv_dataset:
        from datasets.g_buffer_objaverse import load_data, load_eval_data, load_memory_data, load_wds_data
    else:  # shapenet
        from datasets.shapenet import load_data, load_eval_data, load_memory_data

    if args.overfitting:
        data = load_memory_data(
            file_path=args.data_dir,
            batch_size=args.batch_size,
            reso=args.image_size,
            reso_encoder=args.image_size_encoder,  # 224 -> 128
            num_workers=args.num_workers,
            # load_depth=args.depth_lambda > 0
            # load_depth=True,  # for evaluation
                    **args_to_dict(args,
                                   dataset_defaults().keys()))
        eval_data = None
    else:
        if args.use_wds:
            # st()
            if args.data_dir == 'NONE':
                with open(args.shards_lst) as f:
                    shards_lst = [url.strip() for url in f.readlines()]
                data = load_wds_data(
                    shards_lst,  # type: ignore
                    args.image_size,
                    args.image_size_encoder,
                    args.batch_size,
                    args.num_workers,
                    #  plucker_embedding=args.plucker_embedding,
                    #  mv_input=args.mv_input,
                    #  split_chunk_input=args.split_chunk_input,
                    **args_to_dict(args,
                                   dataset_defaults().keys()))

            elif not args.inference:
                data = load_wds_data(args.data_dir,
                                     args.image_size,
                                     args.image_size_encoder,
                                     args.batch_size,
                                     args.num_workers,
                                     plucker_embedding=args.plucker_embedding,
                                     mv_input=args.mv_input,
                                     split_chunk_input=args.split_chunk_input)
            else:
                data = None
            # ! load eval

            if args.eval_data_dir == 'NONE':
                with open(args.eval_shards_lst) as f:
                    eval_shards_lst = [url.strip() for url in f.readlines()]
            else:
                eval_shards_lst = args.eval_data_dir  # auto expanded

            eval_data = load_wds_data(
                eval_shards_lst,  # type: ignore
                args.image_size,
                args.image_size_encoder,
                args.eval_batch_size,
                args.num_workers,
                **args_to_dict(args,
                               dataset_defaults().keys()))
            # load_instance=True) # TODO

        else:

            if args.inference:
                data = None
            else:
                data = load_data(
                    file_path=args.data_dir,
                    batch_size=args.batch_size,
                    reso=args.image_size,
                    reso_encoder=args.image_size_encoder,  # 224 -> 128
                    num_workers=args.num_workers,
                    **args_to_dict(args,
                                   dataset_defaults().keys())
                                   )

                    # load_depth=True # for evaluation

            if args.pose_warm_up_iter > 0:
                overfitting_dataset = load_memory_data(
                    file_path=args.data_dir,
                    batch_size=args.batch_size,
                    reso=args.image_size,
                    reso_encoder=args.image_size_encoder,  # 224 -> 128
                    num_workers=args.num_workers,
                    # load_depth=args.depth_lambda > 0
                    # load_depth=True  # for evaluation
                    **args_to_dict(args,
                                   dataset_defaults().keys()))
                data = [data, overfitting_dataset, args.pose_warm_up_iter]

            eval_data = load_eval_data(
                file_path=args.eval_data_dir,
                batch_size=args.eval_batch_size,
                reso=args.image_size,
                reso_encoder=args.image_size_encoder,  # 224 -> 128
                num_workers=args.num_workers,
                load_depth=True,  # for evaluation
                preprocess=auto_encoder.preprocess,
                wds_split=args.wds_split,
                # interval=args.interval,
                # use_lmdb=args.use_lmdb,
                # plucker_embedding=args.plucker_embedding,
                # load_real=args.load_real,
                # four_view_for_latent=args.four_view_for_latent,
                # load_extra_36_view=args.load_extra_36_view,
                # shuffle_across_cls=args.shuffle_across_cls,
                **args_to_dict(args,
                               dataset_defaults().keys()))

    logger.log("creating data loader done...")

    args.img_size = [args.image_size_encoder]
    # try dry run
    # batch = next(data)
    # batch = None

    # logger.log("creating model and diffusion...")

    # let all processes sync up before starting with a new epoch of training
    dist_util.synchronize()

    # schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion)

    opt = dnnlib.EasyDict(args_to_dict(args, loss_defaults().keys()))
    # opt.max_depth, opt.min_depth = args.rendering_kwargs.ray_end, args.rendering_kwargs.ray_start
    if 'disc' in args.trainer_name:
        loss_class = E3DGE_with_AdvLoss(
            device,
            opt,
            # disc_weight=args.patchgan_disc, # rec_cvD_lambda
            disc_factor=args.patchgan_disc_factor,  # reduce D update speed
            disc_weight=args.patchgan_disc_g_weight).to(device)
    else:
        loss_class = E3DGELossClass(device, opt).to(device)

    # writer = SummaryWriter() # TODO, add log dir

    logger.log("training...")

    TrainLoop = {
        'input_rec': TrainLoop3DRec,
        'nv_rec': TrainLoop3DRecNV,
        # 'nv_rec_patch': TrainLoop3DRecNVPatch,
        'nv_rec_patch': TrainLoop3DRecNVPatchSingleForward,
        'nv_rec_patch_mvE': TrainLoop3DRecNVPatchSingleForwardMV,
        'nv_rec_patch_mvE_disc': TrainLoop3DRecNVPatchSingleForwardMVAdvLoss,
        'nv_rec_patch_mvE_gs': TrainLoop3DRecNVPatchSingleForwardMV_NoCrop,
        'nv_rec_patch_mvE_gs_disc': TrainLoop3DRecNVPatchSingleForwardMV_NoCrop_adv,
    }[args.trainer_name]

    logger.log("creating TrainLoop done...")

    # th._dynamo.config.verbose=True # th212 required
    # th._dynamo.config.suppress_errors = True
    auto_encoder.decoder.rendering_kwargs = args.rendering_kwargs
    train_loop = TrainLoop(
        rec_model=auto_encoder,
        loss_class=loss_class,
        data=data,
        eval_data=eval_data,
        #   compile=args.compile,
        **vars(args))
    # train_loop.rendering_kwargs = args.rendering_kwargs

    if args.inference:
        camera = th.load('eval_pose.pt', map_location=dist_util.dev())
        train_loop.eval_novelview_loop(camera=camera,
                                       save_latent=args.save_latent)
    else:
        train_loop.run_loop()


def create_argparser(**kwargs):
    # defaults.update(model_and_diffusion_defaults())

    defaults = dict(
        seed=0,
        dataset_size=-1,
        trainer_name='input_rec',
        use_amp=False,
        overfitting=False,
        num_workers=4,
        image_size=128,
        image_size_encoder=224,
        iterations=150000,
        anneal_lr=False,
        lr=5e-5,
        weight_decay=0.0,
        lr_anneal_steps=0,
        batch_size=1,
        eval_batch_size=12,
        microbatch=-1,  # -1 disables microbatches
        ema_rate="0.9999",  # comma-separated list of EMA values
        log_interval=50,
        eval_interval=2500,
        save_interval=10000,
        resume_checkpoint="",
        use_fp16=False,
        fp16_scale_growth=1e-3,
        data_dir="",
        eval_data_dir="",
        # load_depth=False, # TODO
        logdir="/mnt/lustre/yslan/logs/nips23/",
        # test warm up pose sampling training
        pose_warm_up_iter=-1,
        inference=False,
        export_latent=False,
        save_latent=False,
        wds_split=1,  # out of 4
    )

    defaults.update(dataset_defaults())  # type: ignore
    defaults.update(encoder_and_nsr_defaults())  # type: ignore
    defaults.update(loss_defaults())

    parser = argparse.ArgumentParser()
    add_dict_to_argparser(parser, defaults)

    return parser


if __name__ == "__main__":
    th.multiprocessing.set_start_method('spawn')

    # os.environ[
    # "TORCH_DISTRIBUTED_DEBUG"] = "DETAIL"  # set to DETAIL for runtime logging.
    # os.environ["TORCH_CPP_LOG_LEVEL"]="INFO"
    # os.environ["NCCL_DEBUG"]="INFO"

    args = create_argparser().parse_args()
    args.local_rank = int(os.environ["LOCAL_RANK"])
    # if os.environ['WORLD_SIZE'] > 1:

    # for multi-node training
    if dist_util.get_world_size() > 1:
        args.global_rank = int(os.environ["RANK"])
    else:
        args.global_rank = 0

    args.gpus = th.cuda.device_count()

    opts = args

    args.rendering_kwargs = rendering_options_defaults(opts)

    # print(args)
    with open(os.path.join(args.logdir, 'args.json'), 'w') as f:
        json.dump(vars(args), f, indent=2)

    # Launch processes.
    print('Launching processes...')

    try:
        training_loop(args)
    # except KeyboardInterrupt as e:
    except Exception as e:
        # print(e)
        traceback.print_exc()
        dist_util.cleanup()  # clean port and socket when ctrl+c