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()