debug.py

import argparse
import datetime
import hashlib
import os
import shutil
import time
import warnings
from collections import OrderedDict
from functools import partial
import glob
import json

os.environ["WANDB_MODE"] = "disabled"

import cv2
import numpy as np
import torch
import torch.cuda.amp as amp
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim
import webdataset as wds
from loguru import logger
from torch.optim.lr_scheduler import MultiStepLR, StepLR, CosineAnnealingLR, OneCycleLR

import utils.config as config

from engine.crog_engine_withoutinsmask import train_with_grasp, validate_with_grasp, validate_without_grasp
from model import build_net
from utils.misc import init_random_seed, set_random_seed, setup_logger
from utils.dataset_new import GraspTransformsGAPlusPlus, GraspAnythingPPWebDataset

warnings.filterwarnings("ignore")
cv2.setNumThreads(0)


def get_parser():
    parser = argparse.ArgumentParser(
        description="Pytorch Referring Expression Segmentation with WebDataset (debug version, no wandb)"
    )
    parser.add_argument(
        "--config",
        default="path to xxx.yaml",
        type=str,
        help="config file"
    )
    parser.add_argument(
        "--opts",
        default=None,
        nargs=argparse.REMAINDER,
        help="override some settings in the config."
    )

    parser.add_argument(
        "--debug_dataset",
        action="store_true",
        help="preview raw dataset samples and log image/text/grasp info"
    )
    parser.add_argument(
        "--debug_dataset_samples",
        default=3,
        type=int,
        help="number of raw dataset samples to preview"
    )
    parser.add_argument(
        "--debug_batch",
        action="store_true",
        help="preview first collated batch"
    )
    parser.add_argument(
        "--debug_model",
        action="store_true",
        help="log model parameter names/shapes/requires_grad"
    )
    parser.add_argument(
        "--debug_log_interval",
        default=50,
        type=int,
        help="reserved for engine-side logging interval"
    )
    parser.add_argument(
        "--debug_save_sample_json",
        action="store_true",
        help="save previewed sample meta to json under output_dir"
    )

    args = parser.parse_args()
    assert args.config is not None
    cfg = config.load_cfg_from_cfg_file(args.config)
    if args.opts is not None:
        cfg = config.merge_cfg_from_list(cfg, args.opts)

    cfg.debug_dataset = args.debug_dataset
    cfg.debug_dataset_samples = args.debug_dataset_samples
    cfg.debug_batch = args.debug_batch
    cfg.debug_model = args.debug_model
    cfg.debug_log_interval = args.debug_log_interval
    cfg.debug_save_sample_json = args.debug_save_sample_json

    return cfg


def build_optimizer(param_list, args):
    optimizer_type = getattr(args, "optimizer_type", "Adam")

    if optimizer_type == "Adam":
        optimizer = torch.optim.Adam(
            param_list,
            lr=args.base_lr,
            weight_decay=args.weight_decay,
            betas=getattr(args, "betas", (0.9, 0.999)),
            eps=getattr(args, "eps", 1e-8),
        )
    elif optimizer_type == "AdamW":
        optimizer = torch.optim.AdamW(
            param_list,
            lr=args.base_lr,
            weight_decay=args.weight_decay,
            betas=getattr(args, "betas", (0.9, 0.999)),
            eps=getattr(args, "eps", 1e-8),
        )
    elif optimizer_type == "SGD":
        optimizer = torch.optim.SGD(
            param_list,
            lr=args.base_lr,
            momentum=getattr(args, "momentum", 0.9),
            weight_decay=args.weight_decay,
            nesterov=getattr(args, "nesterov", True),
        )
    else:
        logger.warning(f"Unknown optimizer type: {optimizer_type}. Defaulting to Adam.")
        optimizer = torch.optim.Adam(
            param_list,
            lr=args.base_lr,
            weight_decay=args.weight_decay,
        )

    return optimizer


def build_scheduler(optimizer, args, data_loader):
    scheduler_type = getattr(args, "scheduler_type", "MultiStepLR")

    if scheduler_type == "MultiStepLR":
        scheduler = MultiStepLR(
            optimizer,
            milestones=args.milestones,
            gamma=args.lr_decay
        )
    elif scheduler_type == "StepLR":
        scheduler = StepLR(
            optimizer,
            step_size=getattr(args, "step_size", 30),
            gamma=getattr(args, "lr_decay", 0.1)
        )
    elif scheduler_type == "CosineAnnealingLR":
        scheduler = CosineAnnealingLR(
            optimizer,
            T_max=args.epochs,
            eta_min=getattr(args, "min_lr", 0)
        )
    elif scheduler_type == "OneCycleLR":
        steps_per_epoch = getattr(args, "steps_per_epoch", None)
        if steps_per_epoch is None:
            try:
                steps_per_epoch = len(data_loader)
            except TypeError:
                raise ValueError(
                    "For WebDataset/IterableDataset, please set args.steps_per_epoch explicitly "
                    "or ensure train_loader has a valid __len__."
                )

        scheduler = OneCycleLR(
            optimizer,
            max_lr=args.base_lr,
            epochs=args.epochs,
            steps_per_epoch=steps_per_epoch,
            pct_start=getattr(args, "pct_start", 0.3),
            anneal_strategy=getattr(args, "anneal_strategy", "cos"),
        )
    elif scheduler_type == "CosineAnnealingWarmRestarts":
        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
            optimizer,
            T_0=getattr(args, "T_0", 10),
            T_mult=getattr(args, "T_mult", 2),
            eta_min=getattr(args, "min_lr", 0),
        )
    else:
        scheduler = MultiStepLR(
            optimizer,
            milestones=args.milestones,
            gamma=args.lr_decay
        )

    warmup_steps = getattr(args, "warmup_steps", 0)
    if warmup_steps > 0:
        from utils.warmup_scheduler import WarmupLRScheduler

        scheduler = WarmupLRScheduler(
            optimizer,
            scheduler,
            warmup_steps=warmup_steps,
            warmup_lr=getattr(args, "warmup_lr", 1e-6),
            warmup_method=getattr(args, "warmup_method", "linear"),
        )
        scheduler._scheduler_type = scheduler_type + "_with_warmup"
        scheduler._warmup_steps = warmup_steps
        scheduler._base_scheduler_type = scheduler_type

    return scheduler


def _stable_bucket_from_string(x: str, num_buckets: int = 100):
    h = hashlib.md5(x.encode("utf-8")).hexdigest()
    return int(h[:8], 16) % num_buckets


def _is_val_sample(sample, val_percent: int):
    sample_id = str(sample.get("sample_id", ""))

    if sample_id == "":
        sample_id = "{}|{}|{}|{}|{}".format(
            sample.get("img_path", ""),
            sample.get("grasp_path", ""),
            sample.get("sentence", ""),
            sample.get("obj_id", ""),
            sample.get("part_id", ""),
        )

    bucket = _stable_bucket_from_string(sample_id, 100)
    return bucket < val_percent


def split_filter_fn(sample, is_train: bool, val_percent: int):
    is_val_sample = _is_val_sample(sample, val_percent)
    return (not is_val_sample) if is_train else is_val_sample


def _safe_shape(x):
    try:
        return tuple(x.shape)
    except Exception:
        return None


def _safe_numel(x):
    try:
        return x.numel()
    except Exception:
        return None


def _safe_minmax(x):
    try:
        if torch.is_tensor(x):
            if x.numel() == 0:
                return "empty"
            xf = x.float()
            return (
                f"min={xf.min().item():.6f}, "
                f"max={xf.max().item():.6f}, "
                f"mean={xf.mean().item():.6f}"
            )
        if isinstance(x, np.ndarray):
            if x.size == 0:
                return "empty"
            xf = x.astype(np.float32)
            return (
                f"min={xf.min():.6f}, "
                f"max={xf.max():.6f}, "
                f"mean={xf.mean():.6f}"
            )
        return None
    except Exception as e:
        return f"minmax_failed: {e}"


def _safe_nan_check(x):
    try:
        if torch.is_tensor(x):
            if not x.is_floating_point():
                return {"has_nan": False, "has_inf": False}
            return {
                "has_nan": bool(torch.isnan(x).any().item()),
                "has_inf": bool(torch.isinf(x).any().item()),
            }
        if isinstance(x, np.ndarray):
            if not np.issubdtype(x.dtype, np.floating):
                return {"has_nan": False, "has_inf": False}
            return {
                "has_nan": bool(np.isnan(x).any()),
                "has_inf": bool(np.isinf(x).any()),
            }
        return None
    except Exception as e:
        return {"nan_check_error": str(e)}


def _basename(x):
    if x is None:
        return None
    try:
        return os.path.basename(str(x))
    except Exception:
        return str(x)


def _short_repr(v, max_len=200):
    s = str(v)
    if len(s) > max_len:
        s = s[:max_len] + " ..."
    return s


def _array_stats(arr):
    try:
        arr = np.asarray(arr)
        stats = {
            "shape": list(arr.shape),
            "dtype": str(arr.dtype),
            "numel": int(arr.size),
            "min": float(arr.min()) if arr.size > 0 else None,
            "max": float(arr.max()) if arr.size > 0 else None,
            "mean": float(arr.mean()) if arr.size > 0 else None,
            "sum": float(arr.sum()) if arr.size > 0 else None,
            "nonzero": int(np.count_nonzero(arr)),
        }
        uniq = np.unique(arr)
        if uniq.size <= 20:
            stats["unique"] = uniq.tolist()
        else:
            stats["unique_head"] = uniq[:20].tolist()
            stats["unique_count"] = int(uniq.size)
        return stats
    except Exception as e:
        return {"stats_error": str(e)}


def _tensor_stats(t):
    try:
        arr = t.detach().cpu().numpy()
        return _array_stats(arr)
    except Exception as e:
        return {"stats_error": str(e)}


def log_mask_stats(mask_dict, prefix="MASK"):
    try:
        if not isinstance(mask_dict, dict):
            logger.info(f"[{prefix}] not dict, type={type(mask_dict)}")
            return

        logger.info(f"[{prefix}] keys={list(mask_dict.keys())}")
        for k, v in mask_dict.items():
            if torch.is_tensor(v):
                stats = _tensor_stats(v)
            else:
                stats = _array_stats(v)
            logger.info(f"[{prefix}] {k} stats = {stats}")
    except Exception as e:
        logger.exception(f"[{prefix}] log_mask_stats failed: {e}")


def _to_jsonable(v):
    if isinstance(v, (str, int, float, bool)) or v is None:
        return v
    if isinstance(v, (list, tuple)):
        return [_to_jsonable(x) for x in v]
    if isinstance(v, dict):
        return {str(k): _to_jsonable(val) for k, val in v.items()}
    if torch.is_tensor(v):
        return {
            "shape": list(v.shape),
            "dtype": str(v.dtype),
            "numel": int(v.numel()),
            "stats": _tensor_stats(v),
        }
    if isinstance(v, np.ndarray):
        return {
            "shape": list(v.shape),
            "dtype": str(v.dtype),
            "numel": int(v.size),
            "stats": _array_stats(v),
        }
    return str(v)


def log_raw_sample(sample, prefix="RAW_SAMPLE"):
    try:
        logger.info(f"[{prefix}] keys = {list(sample.keys())}")

        important_keys = ["sample_id", "img_path", "grasp_path", "sentence", "obj_id", "part_id"]
        for k in important_keys:
            if k in sample:
                logger.info(f"[{prefix}] {k} = {_short_repr(sample[k])}")

        if "img_path" in sample:
            logger.info(f"[{prefix}] image_name = {_basename(sample['img_path'])}")
        if "grasp_path" in sample:
            logger.info(f"[{prefix}] grasp_name = {_basename(sample['grasp_path'])}")

        for k, v in sample.items():
            if k == "grasp_masks" and isinstance(v, dict):
                logger.info(f"[{prefix}] grasp_masks: dict keys={list(v.keys())}")
                log_mask_stats(v, prefix=f"{prefix}_GRASP_MASKS")

            elif torch.is_tensor(v):
                logger.info(
                    f"[{prefix}] tensor[{k}] "
                    f"shape={_safe_shape(v)} dtype={v.dtype} numel={_safe_numel(v)} "
                    f"{_safe_minmax(v)} naninfo={_safe_nan_check(v)}"
                )

            elif isinstance(v, np.ndarray):
                logger.info(
                    f"[{prefix}] ndarray[{k}] "
                    f"shape={v.shape} dtype={v.dtype} numel={v.size} "
                    f"{_safe_minmax(v)} naninfo={_safe_nan_check(v)}"
                )
                logger.info(f"[{prefix}] ndarray[{k}] stats={_array_stats(v)}")
                if v.size <= 100:
                    logger.info(f"[{prefix}] ndarray[{k}] value={v}")

            elif isinstance(v, (list, tuple)):
                logger.info(f"[{prefix}] {k}: type={type(v).__name__}, len={len(v)}")
                if len(v) > 0:
                    logger.info(f"[{prefix}] {k}[0] = {_short_repr(v[0])}")

            elif isinstance(v, dict):
                logger.info(f"[{prefix}] {k}: dict keys={list(v.keys())}")

            else:
                logger.info(f"[{prefix}] {k}: type={type(v).__name__}, value={_short_repr(v)}")

    except Exception as e:
        logger.exception(f"[{prefix}] log_raw_sample failed: {e}")


def sample_to_meta_dict(sample):
    meta = {}
    for k, v in sample.items():
        if k == "grasp_masks" and isinstance(v, dict):
            meta[k] = {}
            for mk, mv in v.items():
                if torch.is_tensor(mv):
                    meta[k][mk] = {
                        "shape": list(mv.shape),
                        "dtype": str(mv.dtype),
                        "numel": int(mv.numel()),
                        "stats": _tensor_stats(mv),
                    }
                else:
                    arr = np.asarray(mv)
                    meta[k][mk] = {
                        "shape": list(arr.shape),
                        "dtype": str(arr.dtype),
                        "numel": int(arr.size),
                        "stats": _array_stats(arr),
                    }
        else:
            meta[k] = _to_jsonable(v)

    if "img_path" in sample:
        meta["image_name"] = _basename(sample["img_path"])
    if "grasp_path" in sample:
        meta["grasp_name"] = _basename(sample["grasp_path"])
    return meta


def log_batch_structure(batch, prefix="BATCH"):
    try:
        logger.info(f"[{prefix}] batch_type = {type(batch)}")

        if isinstance(batch, dict):
            logger.info(f"[{prefix}] keys = {list(batch.keys())}")
            for k, v in batch.items():
                if k == "grasp_masks" and isinstance(v, dict):
                    logger.info(f"[{prefix}] grasp_masks: dict keys = {list(v.keys())}")
                    log_mask_stats(v, prefix=f"{prefix}_GRASP_MASKS")

                elif torch.is_tensor(v):
                    logger.info(
                        f"[{prefix}] tensor[{k}] "
                        f"shape={tuple(v.shape)} dtype={v.dtype} numel={v.numel()} "
                        f"{_safe_minmax(v)} naninfo={_safe_nan_check(v)}"
                    )

                elif isinstance(v, np.ndarray):
                    logger.info(
                        f"[{prefix}] ndarray[{k}] "
                        f"shape={v.shape} dtype={v.dtype} numel={v.size} "
                        f"{_safe_minmax(v)} naninfo={_safe_nan_check(v)}"
                    )
                    logger.info(f"[{prefix}] ndarray[{k}] stats={_array_stats(v)}")

                elif isinstance(v, (list, tuple)):
                    logger.info(f"[{prefix}] {k}: type={type(v).__name__}, len={len(v)}")
                    if len(v) > 0:
                        logger.info(f"[{prefix}] {k}[0] = {_short_repr(v[0])}")

                elif isinstance(v, dict):
                    logger.info(f"[{prefix}] {k}: dict keys = {list(v.keys())}")

                else:
                    logger.info(f"[{prefix}] {k}: type={type(v).__name__}, value={_short_repr(v)}")

        elif isinstance(batch, (list, tuple)):
            logger.info(f"[{prefix}] len = {len(batch)}")
            for i, v in enumerate(batch):
                if torch.is_tensor(v):
                    logger.info(
                        f"[{prefix}] idx={i}, shape={tuple(v.shape)}, dtype={v.dtype}, "
                        f"numel={v.numel()}, {_safe_minmax(v)}, naninfo={_safe_nan_check(v)}"
                    )
                elif isinstance(v, np.ndarray):
                    logger.info(
                        f"[{prefix}] idx={i}, shape={v.shape}, dtype={v.dtype}, "
                        f"numel={v.size}, {_safe_minmax(v)}, naninfo={_safe_nan_check(v)}"
                    )
                elif isinstance(v, (list, tuple)):
                    logger.info(f"[{prefix}] idx={i}, type={type(v).__name__}, len={len(v)}")
                    if len(v) > 0:
                        logger.info(f"[{prefix}] idx={i}[0] = {_short_repr(v[0])}")
                else:
                    logger.info(f"[{prefix}] idx={i}, type={type(v).__name__}, value={_short_repr(v)}")
        else:
            logger.info(f"[{prefix}] unsupported batch type: {type(batch)}")
    except Exception as e:
        logger.exception(f"[{prefix}] log_batch_structure failed: {e}")


def log_model_parameters(model, max_lines=300):
    try:
        logger.info("[MODEL_DEBUG] logging model parameter summary ...")
        count = 0
        for name, p in model.named_parameters():
            logger.info(
                f"[MODEL_DEBUG] name={name}, shape={tuple(p.shape)}, "
                f"requires_grad={p.requires_grad}, numel={p.numel()}"
            )
            count += 1
            if count >= max_lines:
                logger.info(f"[MODEL_DEBUG] truncated after {max_lines} parameters")
                break
    except Exception as e:
        logger.exception(f"[MODEL_DEBUG] failed: {e}")


def preview_dataset_samples(args, tar_files, transform_grasp, is_train=True):
    try:
        logger.info(f"[DATA_DEBUG] Previewing {'train' if is_train else 'val'} raw samples ...")

        preview_builder = GraspAnythingPPWebDataset(
            tar_pattern=tar_files,
            transform_img=None,
            transform_grasp=transform_grasp,
            input_size=args.input_size,
            word_length=args.word_len,
            with_depth=getattr(args, "use_depth", False),
            with_grasp_masks=getattr(args, "use_grasp_masks", True),
            shardshuffle=False,
            sampleshuffle=0,
            repeat=False,
            nodesplitter=wds.split_by_node,
            workersplitter=wds.split_by_worker,
        )

        val_percent = int(getattr(args, "val_percent", 10))
        preview_dataset = preview_builder.get_dataset()
        preview_dataset = preview_dataset.select(
            partial(split_filter_fn, is_train=is_train, val_percent=val_percent)
        )

        save_json = []
        max_samples = int(getattr(args, "debug_dataset_samples", 3))

        for idx, sample in enumerate(preview_dataset):
            log_raw_sample(sample, prefix=f"{'TRAIN' if is_train else 'VAL'}_RAW_SAMPLE_{idx}")
            if getattr(args, "debug_save_sample_json", False):
                save_json.append(sample_to_meta_dict(sample))
            if idx + 1 >= max_samples:
                break

        if getattr(args, "debug_save_sample_json", False) and len(save_json) > 0:
            os.makedirs(args.output_dir, exist_ok=True)
            out_json = os.path.join(
                args.output_dir,
                f"{'train' if is_train else 'val'}_sample_preview.json"
            )
            with open(out_json, "w", encoding="utf-8") as f:
                json.dump(save_json, f, ensure_ascii=False, indent=2)
            logger.info(f"[DATA_DEBUG] saved preview json to {out_json}")

    except Exception as e:
        logger.exception(f"[DATA_DEBUG] preview raw sample failed: {e}")


def preview_first_batch(args, dataset_builder, is_train=True):
    try:
        logger.info(f"[DATA_DEBUG] Previewing first {'train' if is_train else 'val'} batch ...")

        val_percent = int(getattr(args, "val_percent", 10))
        preview_dataset = dataset_builder.get_dataset().select(
            partial(split_filter_fn, is_train=is_train, val_percent=val_percent)
        )

        preview_loader = torch.utils.data.DataLoader(
            preview_dataset,
            batch_size=1,
            shuffle=False,
            num_workers=0,
            pin_memory=False,
            persistent_workers=False,
            drop_last=False,
            collate_fn=GraspAnythingPPWebDataset.collate_fn,
        )

        first_batch = next(iter(preview_loader))
        log_batch_structure(first_batch, prefix=f"{'TRAIN' if is_train else 'VAL'}_FIRST_BATCH")

    except StopIteration:
        logger.error(f"[DATA_DEBUG] {'train' if is_train else 'val'} preview_loader is empty!")
    except Exception as e:
        logger.exception(f"[DATA_DEBUG] preview first batch failed: {e}")


def build_wds_loader(args, is_train=True):
    tar_pattern = args.train_tar
    tar_files = sorted(glob.glob(tar_pattern))
    if len(tar_files) == 0:
        raise FileNotFoundError(f"No tar files matched pattern: {tar_pattern}")

    batch_size = args.batch_size if is_train else args.batch_size_val
    num_workers = args.workers if is_train else args.workers_val

    transform_grasp = GraspTransformsGAPlusPlus(
        width_factor=getattr(args, "width_factor", 100),
        width=args.input_size,
        height=args.input_size,
    )

    if getattr(args, "debug_dataset", False) and getattr(args, "rank", 0) == 0:
        preview_dataset_samples(args, tar_files, transform_grasp, is_train=is_train)

    dataset_builder = GraspAnythingPPWebDataset(
        tar_pattern=tar_files,
        transform_img=None,
        transform_grasp=transform_grasp,
        input_size=args.input_size,
        word_length=args.word_len,
        with_depth=getattr(args, "use_depth", False),
        with_grasp_masks=getattr(args, "use_grasp_masks", True),
        shardshuffle=is_train,
        sampleshuffle=getattr(args, "sampleshuffle", 1000) if is_train else 0,
        repeat=False,
        nodesplitter=wds.split_by_node,
        workersplitter=wds.split_by_worker,
    )

    dataset = dataset_builder.get_dataset()

    val_percent = int(getattr(args, "val_percent", 10))
    dataset = dataset.select(
        partial(split_filter_fn, is_train=is_train, val_percent=val_percent)
    )

    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=(num_workers > 0),
        drop_last=is_train,
        collate_fn=GraspAnythingPPWebDataset.collate_fn,
    )

    if getattr(args, "debug_batch", False) and getattr(args, "rank", 0) == 0:
        preview_first_batch(args, dataset_builder, is_train=is_train)

    return loader


@logger.catch
def main():
    torch.multiprocessing.set_start_method("spawn", force=True)

    args = get_parser()
    args.manual_seed = init_random_seed(args.manual_seed)
    set_random_seed(args.manual_seed, deterministic=False)

    args.ngpus_per_node = torch.cuda.device_count()
    args.world_size = args.ngpus_per_node * args.world_size

    children = []
    for i in range(args.world_size):
        subproc = mp.Process(target=main_worker, args=(i, args))
        children.append(subproc)
        subproc.start()

    for i in range(args.world_size):
        children[i].join()


def main_worker(gpu, args):
    args.output_dir = os.path.join(args.output_folder, args.exp_name)

    args.gpu = gpu
    args.rank = args.rank * args.ngpus_per_node + gpu
    torch.cuda.set_device(args.gpu)

    setup_logger(
        args.output_dir,
        distributed_rank=args.gpu,
        filename="train.log",
        mode="a"
    )

    logger.info("=" * 80)
    logger.info(f"[START] gpu={gpu}, rank={args.rank}, output_dir={args.output_dir}")
    logger.info("=" * 80)

    dist.init_process_group(
        backend=args.dist_backend,
        init_method=args.dist_url,
        world_size=args.world_size,
        rank=args.rank
    )

    model, param_list = build_net(args)
    if args.sync_bn:
        model = nn.SyncBatchNorm.convert_sync_batchnorm(model)

    logger.info(model)
    logger.info(args)

    if args.rank == 0:
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        logger.info(f"Total parameters: {total_params:,}")
        logger.info(f"Trainable parameters: {trainable_params:,}")
        logger.info(f"Using same train_tar for train/val split: {args.train_tar}")
        logger.info(f"Validation percentage from train_tar: {getattr(args, 'val_percent', 10)}%")
        logger.info(f"Optimizer type: {getattr(args, 'optimizer_type', 'Adam')}")
        logger.info(f"Scheduler type: {getattr(args, 'scheduler_type', 'MultiStepLR')}")

        if getattr(args, "debug_model", False):
            log_model_parameters(model)

    model = nn.parallel.DistributedDataParallel(
        model.cuda(),
        device_ids=[args.gpu],
        find_unused_parameters=True
    )

    args.batch_size = int(args.batch_size / args.ngpus_per_node)
    args.batch_size_val = int(args.batch_size_val / args.ngpus_per_node)
    args.workers = int((args.workers + args.ngpus_per_node - 1) / args.ngpus_per_node)
    args.workers_val = int((args.workers_val + args.ngpus_per_node - 1) / args.ngpus_per_node)

    if args.rank == 0:
        logger.info(
            f"[DATALOADER_CFG] per_gpu batch_size={args.batch_size}, "
            f"batch_size_val={args.batch_size_val}, workers={args.workers}, workers_val={args.workers_val}"
        )

    train_loader = build_wds_loader(args, is_train=True)
    val_loader = build_wds_loader(args, is_train=False)

    optimizer = build_optimizer(param_list, args)
    scheduler = build_scheduler(optimizer, args, train_loader)
    scaler = amp.GradScaler()

    if args.rank == 0:
        logger.info(f"[OPTIMIZER] {optimizer}")
        logger.info(f"[SCHEDULER] {scheduler}")
        logger.info(f"[SCALER] enabled={scaler.is_enabled()}")

    dist.barrier()

    best_IoU = 0.0
    best_j_index = 0.0

    if args.resume:
        if os.path.isfile(args.resume):
            logger.info(f"=> loading checkpoint '{args.resume}'")
            map_location = {f"cuda:{0}": f"cuda:{gpu}"}
            checkpoint = torch.load(args.resume, map_location=map_location)

            args.start_epoch = max(0, checkpoint["epoch"] - 1)
            best_IoU = checkpoint.get("best_iou", 0.0)
            if "best_j_index" in checkpoint:
                best_j_index = checkpoint["best_j_index"]

            state_dict = checkpoint["state_dict"]
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                if k.startswith("module."):
                    name = k[7:]
                else:
                    name = k
                new_state_dict[name] = v

            try:
                model.load_state_dict(checkpoint["state_dict"])
            except Exception:
                model.module.load_state_dict(new_state_dict)

            optimizer.load_state_dict(checkpoint["optimizer"])
            if "scheduler" in checkpoint:
                scheduler.load_state_dict(checkpoint["scheduler"])

            logger.info(f"=> loaded checkpoint '{args.resume}' (epoch {checkpoint['epoch']})")
            logger.info(f"=> resume best_IoU = {best_IoU}")
            logger.info(f"=> resume best_j_index = {best_j_index}")

            del checkpoint
            torch.cuda.empty_cache()
        else:
            raise ValueError(
                f"=> resume failed! no checkpoint found at '{args.resume}'. Please check args.resume again!"
            )

    start_time = time.time()
    for epoch in range(args.start_epoch, args.epochs):
        epoch_start_time = time.time()
        epoch_log = epoch + 1

        if args.rank == 0:
            current_lrs = [pg["lr"] for pg in optimizer.param_groups]
            logger.info(f"[EPOCH_START] epoch={epoch_log}, lr={current_lrs}")

        try:
            train_with_grasp(train_loader, model, optimizer, scheduler, scaler, epoch_log, args)
        except Exception as e:
            logger.exception(f"[TRAIN_ERROR] epoch={epoch_log} failed: {e}")
            raise

        epoch_train_time = time.time()

        try:
            if args.use_grasp_masks:
                iou, prec_dict, j_index = validate_with_grasp(val_loader, model, epoch_log, args)
            else:
                iou, prec_dict, j_index = validate_without_grasp(val_loader, model, epoch_log, args)
        except Exception as e:
            logger.exception(f"[VAL_ERROR] epoch={epoch_log} failed: {e}")
            raise

        epoch_training_time = epoch_train_time - epoch_start_time
        epoch_eval_time = time.time() - epoch_train_time

        if args.rank == 0:
            logger.info(
                f"[EPOCH_RESULT] epoch={epoch_log}, "
                f"iou={iou}, j_index={j_index}, "
                f"train_time={epoch_training_time:.2f}s, eval_time={epoch_eval_time:.2f}s"
            )

        if args.rank == 0:
            lastname = os.path.join(args.output_dir, "last_model.pth")
            torch.save(
                {
                    "epoch": epoch_log,
                    "cur_iou": iou,
                    "best_iou": best_IoU,
                    "best_j_index": best_j_index,
                    "prec": prec_dict,
                    "j_index": j_index,
                    "state_dict": model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "scheduler": scheduler.state_dict(),
                },
                lastname
            )

            if iou >= best_IoU:
                best_IoU = iou
                bestname_iou = os.path.join(args.output_dir, "best_iou_model.pth")
                shutil.copyfile(lastname, bestname_iou)
                logger.info(f"[CHECKPOINT] New best IoU = {best_IoU:.6f}, saved to {bestname_iou}")

            if j_index[0] >= best_j_index:
                best_j_index = j_index[0]
                bestname = os.path.join(args.output_dir, "best_jindex_model.pth")
                shutil.copyfile(lastname, bestname)
                logger.info(f"[CHECKPOINT] New best J-index@1 = {best_j_index:.6f}, saved to {bestname}")

        if isinstance(scheduler, OneCycleLR):
            if args.rank == 0:
                logger.info("[SCHEDULER] OneCycleLR detected, assuming per-batch stepping in engine.")
        elif hasattr(scheduler, "_warmup_steps"):
            if getattr(scheduler, "_warmup_complete", False):
                scheduler.base_scheduler.step(epoch_log)
                if hasattr(scheduler.base_scheduler, "_last_lr"):
                    scheduler._last_lr = scheduler.base_scheduler._last_lr
                if args.rank == 0:
                    logger.info(f"[SCHEDULER] warmup complete, base scheduler stepped at epoch={epoch_log}")
            else:
                if args.rank == 0:
                    logger.info(f"[SCHEDULER] warmup not complete at epoch={epoch_log}")
        else:
            scheduler.step(epoch_log)
            if args.rank == 0:
                logger.info(f"[SCHEDULER] stepped at epoch={epoch_log}")

        torch.cuda.empty_cache()

    time.sleep(2)

    logger.info("* Best IoU={} * ".format(best_IoU))
    logger.info("* Best J-index@1={} * ".format(best_j_index))
    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    logger.info("* Training time {} *".format(total_time_str))

    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()