train_acc.py

import os
import re
import json
import random
import wandb
import torch
import dnnlib
from datetime import datetime
from accelerate import Accelerator
from accelerate.utils import set_seed
from training import training_loop_acc as training_loop_acc
from training.utils import Config, process_arguments

import warnings

warnings.filterwarnings("ignore", "Grad strides do not match bucket view strides")  # False warning printed by PyTorch 1.12.

# ----------------------------------------------------------------------------
# Parse a comma separated list of numbers or ranges and return a list of ints.
# Example: '1,2,5-10' returns [1, 2, 5, 6, 7, 8, 9, 10]


def parse_int_list(s):
    if isinstance(s, list):
        return s
    ranges = []
    range_re = re.compile(r"^(\d+)-(\d+)$")
    for p in s.split(","):
        m = range_re.match(p)
        if m:
            ranges.extend(range(int(m.group(1)), int(m.group(2)) + 1))
        else:
            ranges.append(int(p))
    return ranges


def main():
    # Load configuration
    args = process_arguments(default_conf="configs/training/default.yml", debug_conf="configs/training/debug.yml")
    conf = Config(args)

    # Initialize Accelerator
    accelerator = Accelerator(
        mixed_precision='fp16' if conf['fp16'] else 'no',
        gradient_accumulation_steps=1,
        log_with=None,  # Can add wandb, tensorboard etc.
    )

    # Set seed for reproducibility
    if conf["seed"] is not None:
        set_seed(conf["seed"])
    else:
        # Main process generates a seed
        if accelerator.is_main_process:
            seed = torch.randint(1 << 31, size=[]).item()
        else:
            seed = 0  # dummy value

        # Convert to tensor and gather from all processes
        seed_tensor = torch.tensor(seed, device=accelerator.device)
        gathered_seed = accelerator.gather_for_metrics(seed_tensor)

        if gathered_seed.numel() > 1:
            seed = gathered_seed[0].item()
        else:
            seed = gathered_seed.item()

        set_seed(seed)
        conf.update("seed", int(seed))

    # Initialize wandb
    if accelerator.is_main_process:
        wandb.init(
            config=conf.to_dict(),
            name=conf["name"],
            mode=conf["wandb"],
            project=conf['project_name'],
            entity=conf['wandb_usr'],
            #id=resume_run_id,
        )
        wandb.run.log_code(root=".")

    # Initialize config dict.
    c = dnnlib.EasyDict()
    if conf['dataset_type'].lower() == "hf":
        c.dataset_kwargs = dnnlib.EasyDict(
            class_name="training.dataset_hf.PDEDataset",
            path=conf["data"],
            resolution=conf["resolution"],
            use_labels=conf["cond"],
            xflip=conf["xflip"],
            cache=conf["cache"],
        )
    else:
        c.dataset_kwargs = dnnlib.EasyDict(
        class_name='training.dataset.ImageFolderDataset',
        path=conf["data"],
        resolution=conf["resolution"],
        use_labels=conf["cond"],
        xflip=conf["xflip"],
        cache=conf["cache"],
        transpose=conf["transpose"]
    )

    # ===================
    # VALIDATION DATASET
    # ===================
    c.val_dataset_kwargs = None  # Default: no validation
    if conf.get("val_data"):
        val_path = conf["val_data"]
        if os.path.isdir(val_path):
            if accelerator.is_main_process:
                print(f"Using validation dataset from config: {val_path}")
            if conf['dataset_type'].lower() == "hf":
                c.val_dataset_kwargs = dnnlib.EasyDict(
                    class_name="training.dataset_hf.PDEDataset",
                    path=val_path,
                    resolution=conf["resolution"],
                    use_labels=conf["cond"],
                    xflip=False,  # No augmentation for validation
                    cache=conf.get("val_cache", False),  # Usually don't cache validation
                )
            else:
                c.val_dataset_kwargs = dnnlib.EasyDict(
                    class_name='training.dataset.ImageFolderDataset',
                    path=val_path,
                    resolution=conf["resolution"],
                    use_labels=conf["cond"],
                    xflip=False,  # No augmentation for validation
                    cache=conf.get("val_cache", False),
                    transpose=conf["transpose"]
                )
        else:
            if accelerator.is_main_process:
                print(f"Warning: Validation path {val_path} does not exist, skipping validation")
    
    # Validation frequency settings
    c.val_ticks = conf.get("val_ticks", 10)  # Validate every N ticks (default: 10)
    c.val_num_batches = conf.get("val_num_batches", 50)  # Number of batches for validation
    
    # ===================
    # END VALIDATION DATASET
    # ===================

    c.data_loader_kwargs = dnnlib.EasyDict(num_workers=conf["workers"], pin_memory=True, prefetch_factor=2)
    c.network_kwargs = dnnlib.EasyDict()
    c.loss_kwargs = dnnlib.EasyDict()
    c.optimizer_kwargs = dnnlib.EasyDict(class_name="torch.optim.Adam", lr=conf["lr"], betas=[0.9, 0.999], eps=1e-8)
    c.sampler_kwargs = dnnlib.EasyDict(class_name="training.noise_samplers.RBFKernel", scale=conf["rbf_scale"])

    # Validate dataset options.
    try:
        dataset_obj = dnnlib.util.construct_class_by_name(**c.dataset_kwargs)
        c.dataset_kwargs.resolution = dataset_obj.resolution  # be explicit about dataset resolution
        c.dataset_kwargs.max_size = len(dataset_obj)  # be explicit about dataset size
        if conf["cond"] and not dataset_obj.has_labels:
            raise ValueError("--cond=True requires labels specified in dataset.json")
        del dataset_obj  # conserve memory
    except IOError as err:
        raise ValueError(f"--data: {err}")

    # Validate validation dataset if specified
    if c.val_dataset_kwargs is not None:
        try:
            val_dataset_obj = dnnlib.util.construct_class_by_name(**c.val_dataset_kwargs)
            c.val_dataset_kwargs.resolution = val_dataset_obj.resolution
            c.val_dataset_kwargs.max_size = len(val_dataset_obj)
            if accelerator.is_main_process:
                print(f"Validation dataset size: {len(val_dataset_obj)} samples")
            del val_dataset_obj
        except IOError as err:
            if accelerator.is_main_process:
                print(f"Warning: Could not load validation dataset: {err}")
            c.val_dataset_kwargs = None

    # Network architecture.
    if conf["arch"] == "ddpmpp":
        c.network_kwargs.update(model_type="SongUNet", embedding_type="positional", encoder_type="standard", decoder_type="standard")
        c.network_kwargs.update(channel_mult_noise=1, resample_filter=[1, 1], model_channels=128, channel_mult=[2, 2, 2])
    elif conf["arch"] == "ncsnpp":
        c.network_kwargs.update(model_type="SongUNet", embedding_type="fourier", encoder_type="residual", decoder_type="standard")
        c.network_kwargs.update(channel_mult_noise=2, resample_filter=[1, 3, 3, 1], model_channels=128, channel_mult=[2, 2, 2])
    elif conf["arch"] == "adm":
        c.network_kwargs.update(model_type="DhariwalUNet", model_channels=192, channel_mult=[1, 2, 3, 4])
    elif conf["arch"] == "ddpmpp-uno":
        c.network_kwargs.update(model_type="SongUNO", embedding_type="positional", encoder_type="standard", decoder_type="standard")
        c.network_kwargs.update(channel_mult_noise=1, resample_filter=[1, 1], model_channels=128, channel_mult=[2, 2, 2])
        c.network_kwargs.update(
            cond=conf["cond"],
            attn_resolutions=conf["attn_resolutions"],
            num_blocks=conf["num_blocks"],
            fmult=conf["fmult"],
            rank=conf["rank"],
        )
    elif conf["arch"] == "ddpmpp-ulno":
        c.network_kwargs.update(model_type="SongULNO", embedding_type="positional", encoder_type="standard", decoder_type="standard")
        c.network_kwargs.update(channel_mult_noise=1, resample_filter=[1, 1], model_channels=128, channel_mult=[2, 2, 2])
        c.network_kwargs.update(
            cond=conf["cond"],
            attn_resolutions=conf["attn_resolutions"],
            num_blocks=conf["num_blocks"],
            fmult=conf["fmult"],
            rank=conf["rank"],
            # LocalNO specific parameters
            use_differential=conf.get("use_differential", True),
            conv_padding_mode=conf.get("conv_padding_mode", "periodic"),
            fin_diff_kernel_size=conf.get("fin_diff_kernel_size", 3),
        )
    else:
        raise ValueError(f"Invalid architecture: {conf['arch']}")

    # Preconditioning & loss function.
    if conf["precond"] == "vp":
        c.network_kwargs.class_name = "training.networks.VPPrecond"
        c.loss_kwargs.class_name = "training.loss.VPLoss"
    elif conf["precond"] == "ve":
        c.network_kwargs.class_name = "training.networks.VEPrecond"
        c.loss_kwargs.class_name = "training.loss.VELoss"
    elif conf["precond"] == "edm":
        c.network_kwargs.class_name = "training.networks.EDMPrecond"
        c.loss_kwargs.class_name = "training.loss.EDMLossWithSampler" if conf["arch"] == "ddpmpp-uno" or conf["arch"] == "ddpmpp-ulno" else "training.loss.EDMLoss"
    else:
        raise ValueError(f"Invalid preconditioning: {conf['precond']}")

    # Network options.
    if conf["cbase"] is not None:
        c.network_kwargs.model_channels = conf["cbase"]
    if conf["cres"] is not None:
        c.network_kwargs.channel_mult = conf["cres"]
    c.network_kwargs.update(dropout=conf["dropout"], use_fp16=conf["fp16"])
    if conf["nn_resolution"] is not None:
        assert conf["nn_resolution"] <= conf["resolution"]
        c.network_kwargs.update(img_resolution=conf["nn_resolution"])
    else:
        c.network_kwargs.update(img_resolution=conf["resolution"])

    # Training options.
    c.total_kimg = max(int(conf["duration"] * 1000), 1)
    c.lr_rampup_kimg = int(conf["lr_rampup"] * 1000)
    c.ema_halflife_kimg = int(conf["ema"] * 1000)
    c.update(batch_size=conf["batch"], batch_gpu=conf["batch_gpu"])
    c.update(loss_scaling=conf["ls"], cudnn_benchmark=conf["bench"])
    c.update(kimg_per_tick=conf["tick"], snapshot_ticks=conf["snap"], state_dump_ticks=conf["dump"])
    c.cond = conf["cond"]
    c.seed = conf["seed"]

    # Resume training.
    if conf["resume"]:
        
        match = re.fullmatch(r"training-state-(\d+).pt", os.path.basename(conf["resume"]))
        if not match or not os.path.isfile(conf["resume"]):
            raise ValueError("--resume must point to training-state-*.pt from a previous training run")
        
        #c.resume_pkl = os.path.join(os.path.dirname(conf["resume"]), f"network-snapshot-{match.group(1)}.pkl")
        c.resume_pkl = os.path.join(os.path.dirname(conf["resume"]), f"network-snapshot-{match.group(1)}.pt")
        print(f"Resuming from the ckpt {c.resume_pkl}")

        c.resume_kimg = int(match.group(1))
        c.resume_state_dump = conf["resume"]
    else:
        print(f"Initializing a new run!")

    # Pick output directory.
    if not accelerator.is_main_process:
        c.run_dir = None
    else:
        formatted_time = datetime.fromtimestamp(wandb.run.start_time).strftime("%m%d_%H%M%S")
        desc = f"{formatted_time}-{conf['name']}-{wandb.run.id}"
        c.run_dir = os.path.join(conf["outdir"], desc)
        assert not os.path.exists(c.run_dir)

    # Print options.
    if accelerator.is_main_process:
        print()
        print("Training options:")
        print(json.dumps(c, indent=2))
        print()
        print(f"Output directory:        {c.run_dir}")
        print(f"Dataset path:            {c.dataset_kwargs.path}")
        print(f"Validation dataset:      {c.val_dataset_kwargs.path if c.val_dataset_kwargs else 'None'}")
        print(f"Validation frequency:    Every {c.val_ticks} ticks ({c.val_num_batches} batches)")
        print(f"Class-conditional:       {c.dataset_kwargs.use_labels}")
        print(f"Network architecture:    {conf['arch']}")
        print(f"Preconditioning & loss:  {conf['precond']}")
        print(f"Number of GPUs:          {accelerator.num_processes}")
        print(f"Batch size:              {c.batch_size}")
        print(f"Mixed-precision:         {c.network_kwargs.use_fp16}")
        print()

    # Dry run?
    if conf["dry_run"]:
        if accelerator.is_main_process:
            print("Dry run; exiting.")
        return

    # Create output directory.
    if accelerator.is_main_process:
        print("Creating output directory...")
        os.makedirs(c.run_dir, exist_ok=True)
        with open(os.path.join(c.run_dir, "training_options.json"), "wt") as f:
            json.dump(c, f, indent=2)

    # Pass accelerator to training loop
    c.accelerator = accelerator

    # Train.
    training_loop_acc.training_loop(**c)

    if accelerator.is_main_process:
        wandb.finish()


if __name__ == "__main__":
    main()