| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import shutil |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| if str(ROOT) not in sys.path: |
| sys.path.insert(0, str(ROOT)) |
|
|
| from utils.gpu_utils import ( |
| gpu_mapping_message, |
| print_gpu_diagnostics, |
| resolve_gpu, |
| subprocess_gpu_env, |
| ) |
| from utils.dataset_cache import apply_dataloader_cli_overrides, dataset_runtime_summary, print_dataloader_policy |
|
|
| SKIP_EXIT_CODE = 75 |
| FAILED_SUBPROCESS_EXIT_CODE = 70 |
| FAILED_NO_CHECKPOINT_EXIT_CODE = 71 |
| FAILED_EVAL_EXIT_CODE = 72 |
|
|
| EXISTING_LEGACY = { |
| "bifa": ("BiFA", ["python3", "train_wildfire.py"]), |
| "bit_cd": ("BIT_CD", ["python3", "train_wildfire.py"]), |
| "cdmamba": ("CDMamba", ["python3", "train_wildfire.py"]), |
| "change3d": ("Change3D", ["python3", "train_wildfire.py"]), |
| "changeformer": ("ChangeFormer", ["python3", "train_wildfire.py"]), |
| "dsifn": ("IFNet", ["python3", "train_wildfire.py"]), |
| "ifnet": ("IFNet", ["python3", "train_wildfire.py"]), |
| "rsm_cd": ("RSM-CD/change_detection_mamba", ["python3", "train_wildfire.py"]), |
| "schanger": ("SChanger", ["python3", "train_wildfire.py"]), |
| "siam_nestedunet": ("Siam-NestedUNet", ["python3", "train_wildfire.py"]), |
| "stanet": ("STANet", ["python3", "train_wildfire.py"]), |
| } |
|
|
| JSON_LEGACY_TRAINERS = {"bifa", "cdmamba"} |
| MAIN_CD_TRAINERS = {"bit_cd", "changeformer"} |
|
|
| NEW_REPOS = { |
| "fc_ef": ("model_repos/fully_convolutional_change_detection", "https://github.com/rcdaudt/fully_convolutional_change_detection"), |
| "fc_siam_conc": ("model_repos/fully_convolutional_change_detection", "https://github.com/rcdaudt/fully_convolutional_change_detection"), |
| "fc_siam_diff": ("model_repos/fully_convolutional_change_detection", "https://github.com/rcdaudt/fully_convolutional_change_detection"), |
| "dsifn": ("model_repos/DSIFN", "original DSIFN repo still requires URL confirmation"), |
| "changemamba": ("model_repos/ChangeMamba", "https://github.com/ChenHongruixuan/ChangeMamba"), |
| "elgcnet": ("model_repos/elgcnet", "https://github.com/techmn/elgcnet"), |
| "changer": ("model_repos/open-cd", "https://github.com/likyoo/open-cd"), |
| "hanet": ("model_repos/HANet-CD", "https://github.com/ChengxiHAN/HANet-CD"), |
| "cgnet": ("model_repos/CGNet-CD", "https://github.com/ChengxiHAN/CGNet-CD"), |
| "dsamnet": ("model_repos/DSAMNet", "https://github.com/liumency/DSAMNet"), |
| "tinycd": ("model_repos/Tiny_model_4_CD", "https://github.com/AndreaCodegoni/Tiny_model_4_CD"), |
| } |
|
|
| TORCHVISION_WEIGHTS = { |
| "changer": "resnet18", |
| "dsamnet": "resnet18", |
| "hanet": "resnet50", |
| "cgnet": "resnet50", |
| "dsifn": "vgg16", |
| } |
|
|
| TIMM_WEIGHTS = { |
| "tinycd": "efficientnet_b4", |
| "elgcnet": "mit_b0", |
| "changeformer": "mit_b1", |
| } |
|
|
| EXTERNAL_TRAINABLE = {"dsamnet", "cgnet", "hanet", "tinycd"} |
|
|
|
|
| def parse_args(model_name: str) -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=f"Train/evaluate {model_name} on a cd-models dataset.") |
| parser.add_argument("--dataset", default="wildfire_s2") |
| parser.add_argument("--epochs", type=int, default=None) |
| parser.add_argument("--batch-size", type=int, default=None) |
| parser.add_argument("--lr", type=float, default=None) |
| parser.add_argument("--gpu", default="0") |
| parser.add_argument("--resume", action="store_true") |
| parser.add_argument("--eval-only", action="store_true") |
| parser.add_argument("--dry-run", action="store_true") |
| parser.add_argument("--smoke-test", action="store_true") |
| parser.add_argument("--force", action="store_true") |
| parser.add_argument("--output-dir", default=None) |
| parser.add_argument("--model", default=model_name, help="Variant selector used by FC-Siam wrappers.") |
| parser.add_argument("--max-iters", type=int, default=None, help="Iteration override for iteration-based upstream trainers.") |
| parser.add_argument("--num-workers", type=int, default=None) |
| parser.add_argument("--prefetch-factor", type=int, default=None) |
| parser.set_defaults(persistent_workers=None, pin_memory=None) |
| parser.add_argument("--persistent-workers", dest="persistent_workers", action="store_true") |
| parser.add_argument("--no-persistent-workers", dest="persistent_workers", action="store_false") |
| parser.add_argument("--pin-memory", dest="pin_memory", action="store_true") |
| parser.add_argument("--no-pin-memory", dest="pin_memory", action="store_false") |
| return parser.parse_args() |
|
|
|
|
| def _ensure_pretrained(model_name: str, variant: str = "tiny") -> str | None: |
| from utils.weight_downloader import ensure_timm_weight, ensure_torchvision_weight, ensure_weights |
|
|
| if model_name == "changemamba": |
| return ensure_weights(f"vmamba_{variant}") |
| if model_name in TIMM_WEIGHTS: |
| ensure_timm_weight(TIMM_WEIGHTS[model_name], required=model_name == "tinycd") |
| if model_name in TORCHVISION_WEIGHTS: |
| ensure_torchvision_weight(TORCHVISION_WEIGHTS[model_name]) |
| return None |
|
|
|
|
| def _dataset_ready(dataset_name: str, return_format: str = "tuple", dry_run: bool = False) -> dict: |
| from datasets.cd_dataset import CDDataset |
| from utils.config_loader import load_dataset_config |
|
|
| cfg = load_dataset_config(dataset_name) |
| root_exists = Path(cfg["data_root"]).exists() |
| print( |
| f"[DATASET] {cfg['name']} root={cfg['data_root']} img_size={cfg.get('img_size')} " |
| f"channels={cfg.get('channels')} return_format={return_format}" |
| ) |
| print(f"[DATASET] {dataset_runtime_summary(cfg)}") |
| print_dataloader_policy(cfg, torch.cuda.is_available()) |
| if cfg.get("io_warning"): |
| print(f"[DATASET-WARNING] {cfg['io_warning']}") |
| print(f"[DRY-RUN] Dataset root exists: {root_exists} -> {cfg['data_root']}" if dry_run else f"[DATASET] Root exists: {root_exists}") |
| if dry_run: |
| return cfg |
| if not root_exists: |
| raise FileNotFoundError( |
| f"Dataset root directory not found for {dataset_name}: {cfg['data_root']}. " |
| "Set DATA_ROOT or edit the dataset YAML." |
| ) |
| try: |
| from utils.dataset_list_generator import generate_list_files |
|
|
| list_dir = Path(cfg["data_root"]) / "list" |
| if not (list_dir / "train.txt").exists(): |
| print(f"[SETUP] Generating list files for {dataset_name}...") |
| generate_list_files( |
| cfg["data_root"], |
| splits=list((cfg.get("splits") or {"train": "train", "val": "val", "test": "test"}).keys()), |
| img_subdir=cfg.get("image_a_folder", "A"), |
| ) |
| except Exception as exc: |
| print(f"[SETUP] List-file generation skipped: {exc}") |
| if not Path(cfg["data_root"]).exists(): |
| print(f"[DATASET] Root does not exist yet, skipping CDDataset construction: {cfg['data_root']}") |
| return cfg |
| _ = CDDataset(cfg["data_root"], "train", cfg=cfg, return_format=return_format) |
| return cfg |
|
|
|
|
| def _run_logged(cmd: list[str], cwd: Path, env: dict, stdout_path: Path, stderr_path: Path) -> int: |
| stdout_path.parent.mkdir(parents=True, exist_ok=True) |
| stderr_path.parent.mkdir(parents=True, exist_ok=True) |
| visible = env.get("CUDA_VISIBLE_DEVICES") |
| if visible: |
| print(_gpu_mapping_message(env)) |
| with stdout_path.open("w", encoding="utf-8") as stdout, stderr_path.open("w", encoding="utf-8") as stderr: |
| return subprocess.run(cmd, cwd=cwd, env=env, stdout=stdout, stderr=stderr, check=False).returncode |
|
|
|
|
| def _run_streamed(cmd: list[str], cwd: Path, env: dict, log_path: Path) -> int: |
| log_path.parent.mkdir(parents=True, exist_ok=True) |
| with log_path.open("a", encoding="utf-8") as log: |
| log.write("\n[WRAPPER] cwd=" + str(cwd) + "\n") |
| log.write("[WRAPPER] " + _gpu_mapping_message(env) + "\n") |
| log.write("[WRAPPER] command=" + " ".join(cmd) + "\n") |
| log.flush() |
| process = subprocess.Popen( |
| cmd, |
| cwd=cwd, |
| env=env, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT, |
| text=True, |
| bufsize=1, |
| ) |
| assert process.stdout is not None |
| for line in process.stdout: |
| print(line, end="", flush=True) |
| log.write(line) |
| log.flush() |
| return process.wait() |
|
|
|
|
| def _gpu_mapping_message(env: dict) -> str: |
| return gpu_mapping_message(env=env) |
|
|
|
|
| def _print_train_command(cwd: Path, cmd: list[str], env: dict) -> None: |
| print(_gpu_mapping_message(env)) |
| print("[TRAIN]", cwd, " ".join(cmd)) |
|
|
|
|
| def _evaluate(model_name: str, dataset: str, gpu: str, dry_run: bool) -> int: |
| cmd = [sys.executable, str(ROOT / "evaluate.py"), "--model", model_name, "--dataset", dataset, "--gpu", gpu] |
| print("[EVAL]", " ".join(cmd)) |
| if dry_run: |
| return 0 |
| log_dir = ROOT / "results" / model_name / dataset / "logs" |
| env, _ = subprocess_gpu_env(gpu) |
| return _run_logged(cmd, ROOT, env, log_dir / "eval_stdout.log", log_dir / "eval_stderr.log") |
|
|
|
|
| def _legacy_checkpoint_candidates(model_name: str, dataset_name: str) -> list[Path]: |
| if model_name != "cdmamba": |
| return [] |
|
|
| experiments = ROOT / "CDMamba" / "experiments" |
| if not experiments.exists(): |
| return [] |
|
|
| candidates: list[Path] = [] |
| run_dirs = sorted( |
| ( |
| path |
| for path in experiments.glob(f"{dataset_name}-train-cdmamba_*") |
| if path.is_dir() |
| ), |
| key=lambda path: path.stat().st_mtime, |
| ) |
| preferred_patterns = [ |
| "checkpoint/best_cd_model_gen.pth", |
| "checkpoint/*_gen.pth", |
| "checkpoint/*best*.pth", |
| "checkpoint/*.pth", |
| "**/*.pth", |
| "**/*.pt", |
| ] |
| for run_dir in run_dirs: |
| for pattern in preferred_patterns: |
| matches = [ |
| path |
| for path in sorted(run_dir.glob(pattern), key=lambda item: item.stat().st_mtime) |
| if path.is_file() and "_opt" not in path.name.lower() |
| ] |
| if matches: |
| candidates.extend(matches) |
| break |
| return candidates |
|
|
|
|
| def _canonicalize_checkpoints(model_name: str, dataset_name: str) -> None: |
| out_dir = ROOT / "results" / model_name / dataset_name |
| ckpt_dir = out_dir / "checkpoints" |
| ckpt_dir.mkdir(parents=True, exist_ok=True) |
| best_out = ckpt_dir / "best_model.pth" |
| latest_out = ckpt_dir / "latest.pth" |
| if best_out.exists(): |
| return |
| patterns = [ |
| "**/best_model.pth", |
| "**/best_ckpt.pth", |
| "**/best_ckpt.pt", |
| "**/*best*.pth", |
| "**/*best*.pt", |
| "**/*_best_iou.pth", |
| "**/netCD_epoch_*.pth", |
| "**/*.pth", |
| "**/*.pt", |
| ] |
| candidates: list[Path] = [] |
| for pattern in patterns: |
| for path in sorted(out_dir.glob(pattern)): |
| if path.is_file() and path.resolve() != best_out.resolve(): |
| candidates.append(path) |
| if candidates: |
| break |
| if not candidates: |
| candidates = _legacy_checkpoint_candidates(model_name, dataset_name) |
| if not candidates: |
| raise FileNotFoundError( |
| f"Training completed for {model_name}/{dataset_name}, but no checkpoint file was found under " |
| f"{out_dir} or known legacy output directories." |
| ) |
| source = candidates[-1] |
| shutil.copy2(source, best_out) |
| if not latest_out.exists(): |
| shutil.copy2(source, latest_out) |
| metadata = { |
| "model": model_name, |
| "dataset": dataset_name, |
| "canonical_best_model": str(best_out), |
| "source_checkpoint": str(source), |
| "source_checkpoint_name": source.name, |
| } |
| (ckpt_dir / "checkpoint_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") |
|
|
|
|
| def _checkpoint_files(checkpoint_dir: Path) -> list[Path]: |
| suffixes = {".pth", ".pt", ".ckpt"} |
| if not checkpoint_dir.exists(): |
| return [] |
| return sorted(path for path in checkpoint_dir.rglob("*") if path.is_file() and path.suffix.lower() in suffixes) |
|
|
|
|
| def _print_checkpoint_failure(model_name: str, dataset_name: str, cmd: list[str], checkpoint_dir: Path, result_dir: Path) -> None: |
| print(f"[FAILED_NO_CHECKPOINT] {model_name}/{dataset_name}: upstream command returned successfully but no checkpoint was found.") |
| print(f"[FAILED_NO_CHECKPOINT] command={' '.join(cmd)}") |
| print(f"[FAILED_NO_CHECKPOINT] expected checkpoint dir={checkpoint_dir}") |
| if result_dir.exists(): |
| existing = [str(path.relative_to(result_dir)) for path in sorted(result_dir.rglob("*"))[:80]] |
| print(f"[FAILED_NO_CHECKPOINT] existing result files={existing}") |
| else: |
| print(f"[FAILED_NO_CHECKPOINT] result dir does not exist: {result_dir}") |
|
|
|
|
| def _iteration_budget(model_name: str, model_cfg: dict, args: argparse.Namespace, default: int) -> tuple[int, str]: |
| if args.max_iters is not None: |
| return int(args.max_iters), "cli --max-iters" |
| if args.smoke_test: |
| return int(model_cfg.get("smoke_max_iters", min(default, 2))), "model smoke_max_iters" |
| if "max_iters" in model_cfg: |
| return int(model_cfg["max_iters"]), "model max_iters" |
| return default, f"{model_name} wrapper default" |
|
|
|
|
| def _model_config(model_name: str) -> dict: |
| from utils.config_loader import load_model_config |
|
|
| return load_model_config(model_name) |
|
|
|
|
| def _legacy_dataset_token(dataset_cfg: dict) -> str: |
| name = str(dataset_cfg.get("name", "")).lower() |
| source = str(dataset_cfg.get("source_name", "")).lower() |
| if "dsifn" in name or "dsifn" in source: |
| return "DSIFN" |
| if "levir" in name or "levir" in source: |
| return "LEVIR" |
| if "whu" in name or "whu" in source: |
| return "WHU" |
| if "wildfire" in name or "wildfire" in source: |
| return "WildFireS2" |
| return dataset_cfg.get("source_name", dataset_cfg["name"]) |
|
|
|
|
| def _split_dir(dataset_cfg: dict, split: str) -> Path: |
| return Path(dataset_cfg["data_root"]) / dataset_cfg.get("splits", {}).get(split, split) |
|
|
|
|
| def _safe_link_or_copy(src: Path, dst: Path) -> None: |
| if dst.is_symlink(): |
| try: |
| if Path(os.readlink(dst)) == src: |
| return |
| except OSError: |
| pass |
| dst.unlink() |
| if dst.exists(): |
| return |
| dst.parent.mkdir(parents=True, exist_ok=True) |
| if not src.exists(): |
| print(f"[DATASET-VIEW] Source path is not present yet, skipping generated link: {src}") |
| return |
| try: |
| os.symlink(src, dst, target_is_directory=src.is_dir()) |
| except OSError: |
| if src.is_dir(): |
| shutil.copytree(src, dst, dirs_exist_ok=True) |
| else: |
| shutil.copy2(src, dst) |
|
|
|
|
| def _safe_mask_link_or_copy(src: Path, dst: Path) -> None: |
| dst.parent.mkdir(parents=True, exist_ok=True) |
| if not src.exists(): |
| print(f"[DATASET-VIEW] Source mask is not present yet, skipping generated link: {src}") |
| return |
|
|
| try: |
| import numpy as np |
| from PIL import Image |
|
|
| arr = np.asarray(Image.open(src)) |
| if arr.ndim == 3: |
| arr = arr[..., 0] |
| if arr.size and int(arr.max()) <= 1: |
| if dst.exists() or dst.is_symlink(): |
| dst.unlink() |
| Image.fromarray((arr > 0).astype(np.uint8) * 255).save(dst, format="PNG") |
| return |
| except Exception as exc: |
| print(f"[DATASET-VIEW] Could not inspect mask {src}: {exc}") |
|
|
| _safe_link_or_copy(src, dst) |
|
|
|
|
| def _scan_images(folder: Path) -> dict[str, Path]: |
| exts = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"} |
| if not folder.is_dir(): |
| return {} |
| return { |
| p.stem: p |
| for p in sorted(folder.iterdir()) |
| if p.is_file() and p.suffix.lower() in exts |
| } |
|
|
|
|
| def _view_split_counts(dataset_cfg: dict) -> dict[str, int]: |
| counts: dict[str, int] = {} |
| for split in ("train", "val", "test"): |
| split_root = _split_dir(dataset_cfg, split) |
| a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) |
| b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) |
| label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) |
| counts[split] = len(set(a_files) & set(b_files) & set(label_files)) |
| return counts |
|
|
|
|
| def _log_generated_view(dataset_cfg: dict, view: Path) -> None: |
| counts = _view_split_counts(dataset_cfg) |
| print(f"[DATASET] dataset name: {dataset_cfg['name']}") |
| print(f"[DATASET] original root: {dataset_cfg.get('original_data_root', dataset_cfg.get('data_root'))}") |
| print(f"[DATASET] local root: {dataset_cfg.get('local_root')}") |
| print(f"[DATASET] using local root: {bool(dataset_cfg.get('using_local_root'))}") |
| print(f"[DATASET] generated view path: {view}") |
| print(f"[DATASET] train/val/test counts: {counts}") |
|
|
|
|
| def _prepare_matched_split_view(dataset_cfg: dict, split: str) -> Path: |
| split_root = _split_dir(dataset_cfg, split) |
| view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "matched" / split |
| a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) |
| b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) |
| label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) |
| for folder in ("A", "B", "label"): |
| (view / folder).mkdir(parents=True, exist_ok=True) |
| for stem in sorted(set(a_files) & set(b_files) & set(label_files)): |
| target_name = a_files[stem].name |
| _safe_link_or_copy(a_files[stem], view / "A" / target_name) |
| _safe_link_or_copy(b_files[stem], view / "B" / target_name) |
| _safe_mask_link_or_copy(label_files[stem], view / "label" / target_name) |
| _log_generated_view(dataset_cfg, view) |
| return view |
|
|
|
|
| def _prepare_split_view(dataset_cfg: dict, split: str) -> Path: |
| split_root = _split_dir(dataset_cfg, split) |
| view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / split |
| mapping = { |
| "A": split_root / dataset_cfg.get("image_a_folder", "A"), |
| "B": split_root / dataset_cfg.get("image_b_folder", "B"), |
| "label": split_root / dataset_cfg.get("mask_folder", "label"), |
| } |
| for target, src in mapping.items(): |
| _safe_link_or_copy(src, view / target) |
| _log_generated_view(dataset_cfg, view) |
| return view |
|
|
|
|
| def _prepare_tinycd_view(dataset_cfg: dict) -> Path: |
| view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "tinycd" |
| list_dir = view / "list" |
| list_dir.mkdir(parents=True, exist_ok=True) |
| for folder in ("A", "B", "label"): |
| (view / folder).mkdir(parents=True, exist_ok=True) |
| for split in ("train", "val", "test"): |
| split_root = _split_dir(dataset_cfg, split) |
| a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) |
| b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) |
| label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) |
| names = [] |
| for stem in sorted(set(a_files) & set(b_files) & set(label_files)): |
| target_name = f"{split}__{a_files[stem].name}" |
| _safe_link_or_copy(a_files[stem], view / "A" / target_name) |
| _safe_link_or_copy(b_files[stem], view / "B" / target_name) |
| _safe_mask_link_or_copy(label_files[stem], view / "label" / target_name) |
| names.append(target_name) |
| (list_dir / f"{split}.txt").write_text("\n".join(names) + ("\n" if names else ""), encoding="utf-8") |
| _log_generated_view(dataset_cfg, view) |
| return view |
|
|
|
|
| def _prepare_changemamba_view(dataset_cfg: dict) -> Path: |
| view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "changemamba" |
| split_map = {"train": "train", "test": "val"} |
| for target_split, source_split in split_map.items(): |
| for folder in ("T1", "T2", "GT"): |
| (view / target_split / folder).mkdir(parents=True, exist_ok=True) |
| split_root = _split_dir(dataset_cfg, source_split) |
| a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) |
| b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) |
| label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) |
| names = [] |
| for stem in sorted(set(a_files) & set(b_files) & set(label_files)): |
| target_name = f"{source_split}__{a_files[stem].name}" |
| _safe_link_or_copy(a_files[stem], view / target_split / "T1" / target_name) |
| _safe_link_or_copy(b_files[stem], view / target_split / "T2" / target_name) |
| _safe_mask_link_or_copy(label_files[stem], view / target_split / "GT" / target_name) |
| names.append(target_name) |
| list_name = "train_set.txt" if target_split == "train" else "test_set.txt" |
| (view / list_name).write_text("\n".join(names) + ("\n" if names else ""), encoding="utf-8") |
| _log_generated_view(dataset_cfg, view) |
| return view |
|
|
|
|
| def _prepare_opencd_view(dataset_cfg: dict) -> Path: |
| view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "opencd" |
| for split in ("train", "val", "test"): |
| for folder in ("A", "B", "label"): |
| (view / split / folder).mkdir(parents=True, exist_ok=True) |
| split_root = _split_dir(dataset_cfg, split) |
| a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) |
| b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) |
| label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) |
| for stem in sorted(set(a_files) & set(b_files) & set(label_files)): |
| image_name = a_files[stem].name |
| label_name = f"{a_files[stem].stem}{label_files[stem].suffix}" |
| _safe_link_or_copy(a_files[stem], view / split / "A" / image_name) |
| _safe_link_or_copy(b_files[stem], view / split / "B" / image_name) |
| _safe_link_or_copy(label_files[stem], view / split / "label" / label_name) |
| _log_generated_view(dataset_cfg, view) |
| return view |
|
|
|
|
| def _write_opencd_changer_config(dataset_cfg: dict, model_cfg: dict, view: Path) -> Path: |
| train_split = _split_dir(dataset_cfg, "train") |
| n_train = len(_scan_images(train_split / dataset_cfg.get("image_a_folder", "A"))) |
| batch_size = int(dataset_cfg.get("batch_size", 8)) |
| steps_per_epoch = max((n_train + batch_size - 1) // batch_size, 1) |
| max_iters = int(model_cfg.get("max_iters", int(model_cfg.get("num_epochs", 200)) * steps_per_epoch)) |
| val_interval = max(steps_per_epoch, 1) |
| cfg_path = ROOT / "generated_configs" / f"{dataset_cfg['name']}__changer_opencd.py" |
| cfg_path.parent.mkdir(parents=True, exist_ok=True) |
| cfg_path.write_text( |
| "\n".join( |
| [ |
| f"_base_ = '{ROOT / 'model_repos' / 'open-cd' / 'configs' / 'changer' / 'changer_ex_r18_512x512_40k_levircd.py'}'", |
| "", |
| "dataset_type = 'DSIFN_Dataset'", |
| f"data_root = r'{view}'", |
| f"crop_size = ({int(dataset_cfg.get('img_size', model_cfg.get('img_size', 256)))}, {int(dataset_cfg.get('img_size', model_cfg.get('img_size', 256)))})", |
| "", |
| "train_dataloader = dict(", |
| f" batch_size={batch_size},", |
| f" num_workers={int(dataset_cfg.get('num_workers', 4))},", |
| " dataset=dict(", |
| " type=dataset_type,", |
| " data_root=data_root,", |
| " data_prefix=dict(seg_map_path='train/label', img_path_from='train/A', img_path_to='train/B')))", |
| "val_dataloader = dict(", |
| " dataset=dict(", |
| " type=dataset_type,", |
| " data_root=data_root,", |
| " data_prefix=dict(seg_map_path='val/label', img_path_from='val/A', img_path_to='val/B')))", |
| "test_dataloader = dict(", |
| " dataset=dict(", |
| " type=dataset_type,", |
| " data_root=data_root,", |
| " data_prefix=dict(seg_map_path='test/label', img_path_from='test/A', img_path_to='test/B')))", |
| "", |
| f"train_cfg = dict(type='IterBasedTrainLoop', max_iters={max_iters}, val_interval={val_interval})", |
| f"work_dir = r'{ROOT / 'results' / 'changer' / dataset_cfg['name'] / 'work_dir'}'", |
| "", |
| ] |
| ), |
| encoding="utf-8", |
| ) |
| return cfg_path |
|
|
|
|
| def _write_hanet_metadata(dataset_cfg: dict, model_cfg: dict, view_root: Path) -> Path: |
| out_dir = ROOT / "results" / "hanet" / dataset_cfg["name"] |
| metadata = { |
| "patch_size": int(dataset_cfg.get("img_size", 256)), |
| "augmentation": True, |
| "num_gpus": 1, |
| "num_workers": int(dataset_cfg.get("num_workers", 4)), |
| "num_channel": 3, |
| "EF": False, |
| "epochs": int(model_cfg.get("num_epochs", 50)), |
| "epochs_threshold": 15, |
| "gamma": 0.5, |
| "weight_decay": float(model_cfg.get("weight_decay", 5e-4) or 5e-4), |
| "batch_size": int(dataset_cfg.get("batch_size", 8)), |
| "learning_rate": float(model_cfg.get("lr", 5e-4)), |
| "loss_function": "hybrid", |
| "dataset_dir": str(view_root) + "/", |
| "weight_dir": str(out_dir / "weights") + "/", |
| "Output_dir": str(out_dir / "outputs") + "/", |
| "log_dir": str(out_dir / "logs"), |
| } |
| path = ROOT / "generated_configs" / f"{dataset_cfg['name']}__hanet_metadata.json" |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps(metadata, indent=2), encoding="utf-8") |
| return path |
|
|
|
|
| def _existing_train_command(model_name: str, dataset_cfg: dict, args: argparse.Namespace) -> tuple[Path, list[str], dict]: |
| workdir_rel, default_cmd = EXISTING_LEGACY[model_name] |
| workdir = ROOT / workdir_rel |
| env, _ = subprocess_gpu_env(args.gpu) |
| child_gpu = "0" |
| env["CD_MODELS_DATASET_ROOT"] = dataset_cfg["data_root"] |
| env["CD_MODELS_DSIFN_ROOT"] = dataset_cfg["data_root"] |
| env["CD_MODELS_DATASET_NAME"] = dataset_cfg["name"] |
| env["CD_MODELS_NUM_WORKERS"] = str(int(dataset_cfg.get("num_workers", 2))) |
| env["CD_MODELS_PREFETCH_FACTOR"] = str(int(dataset_cfg.get("prefetch_factor", 2))) |
| env["CD_MODELS_PERSISTENT_WORKERS"] = "1" if dataset_cfg.get("persistent_workers") else "0" |
| env["CD_MODELS_PIN_MEMORY"] = "1" if dataset_cfg.get("pin_memory", True) else "0" |
| env["CD_MODELS_BATCH_SIZE"] = str(int(args.batch_size or dataset_cfg.get("batch_size", 8))) |
| if args.epochs is not None: |
| env["CD_MODELS_EPOCHS"] = str(int(args.epochs)) |
| if dataset_cfg["name"] == "wildfire_s2": |
| return workdir, [sys.executable, default_cmd[1]], env |
| if model_name in JSON_LEGACY_TRAINERS: |
| from utils.legacy_config_writer import write_bifa_or_cdmamba_config |
|
|
| model_cfg = _model_config(model_name) |
| if args.epochs is not None: |
| model_cfg["num_epochs"] = int(args.epochs) |
| if args.lr is not None: |
| model_cfg["lr"] = float(args.lr) |
| if args.batch_size is not None: |
| dataset_cfg["batch_size"] = int(args.batch_size) |
| config_path = write_bifa_or_cdmamba_config( |
| model_name, |
| dataset_cfg, |
| model_cfg, |
| prepare_view=not args.dry_run, |
| ) |
| return workdir, [sys.executable, "train_cd.py", "--config", str(config_path), "--phase", "train", "--gpu_ids", child_gpu], env |
| if model_name in MAIN_CD_TRAINERS: |
| from utils.legacy_config_writer import prepare_legacy_list_view |
|
|
| legacy_root = ( |
| ROOT / "generated_dataset_views" / dataset_cfg["name"] / "legacy_list" |
| if args.dry_run |
| else prepare_legacy_list_view(dataset_cfg) |
| ) |
| env["CD_MODELS_DATASET_ROOT"] = str(legacy_root) |
| env["CD_MODELS_DSIFN_ROOT"] = str(legacy_root) |
| model_cfg = _model_config(model_name) |
| token = _legacy_dataset_token(dataset_cfg) |
| project = f"{dataset_cfg['name']}-train-{model_name}" |
| if model_name == "bit_cd": |
| net_g = "base_transformer_pos_s4_dd8" |
| optimizer = "sgd" |
| loss = "ce" |
| lr = float(model_cfg.get("lr", 0.01)) |
| else: |
| net_g = "ChangeFormerV6" |
| optimizer = "adamw" |
| loss = "ce" |
| lr = float(model_cfg.get("lr", 0.00006)) |
| cmd = [ |
| sys.executable, |
| "main_cd.py", |
| "--gpu_ids", |
| child_gpu, |
| "--project_name", |
| project, |
| "--data_name", |
| token, |
| "--img_size", |
| str(int(dataset_cfg.get("img_size", model_cfg.get("img_size", 256)))), |
| "--batch_size", |
| str(int(dataset_cfg.get("batch_size", 8))), |
| "--num_workers", |
| str(int(dataset_cfg.get("num_workers", 4))), |
| "--max_epochs", |
| str(int(model_cfg.get("num_epochs", 200))), |
| "--optimizer", |
| optimizer, |
| "--lr", |
| str(lr), |
| "--loss", |
| loss, |
| "--net_G", |
| net_g, |
| ] |
| return workdir, cmd, env |
| from utils.legacy_config_writer import prepare_legacy_list_view |
|
|
| legacy_root = ( |
| ROOT / "generated_dataset_views" / dataset_cfg["name"] / "legacy_list" |
| if args.dry_run |
| else prepare_legacy_list_view(dataset_cfg) |
| ) |
| out_dir = ROOT / "results" / model_name / dataset_cfg["name"] |
| env["CD_MODELS_DATASET_ROOT"] = str(legacy_root) |
| env["CD_MODELS_DSIFN_ROOT"] = str(legacy_root) |
| env["CD_MODELS_DATASET_NAME"] = dataset_cfg["name"] |
| env["CD_MODELS_CHECKPOINT_DIR"] = str(out_dir / "checkpoints") |
| env["CD_MODELS_LOG_PATH"] = str(out_dir / "logs" / "train.log") |
| if args.lr is not None: |
| env["CD_MODELS_LR"] = str(float(args.lr)) |
| else: |
| env["CD_MODELS_LR"] = str(float(_model_config(model_name).get("lr", 1e-3))) |
| return workdir, [sys.executable, "train_wildfire.py"], env |
|
|
|
|
| def run_existing_model(model_name: str, args: argparse.Namespace) -> int: |
| from utils.model_adapters import get_model_adapter |
| from utils.unified_trainer import train_with_adapter |
|
|
| adapter = get_model_adapter(model_name) |
| if adapter.supports_unified_training: |
| return train_with_adapter(model_name, args) |
| dataset_cfg = _dataset_ready(args.dataset, return_format="legacy", dry_run=args.dry_run) |
| apply_dataloader_cli_overrides(dataset_cfg, args) |
| print_dataloader_policy(dataset_cfg, torch.cuda.is_available()) |
| workdir_rel, cmd = EXISTING_LEGACY[model_name] |
| workdir, cmd, env = _existing_train_command(model_name, dataset_cfg, args) |
| if not cmd: |
| print( |
| f"[SKIP] {model_name}/{args.dataset}: no verified non-WildFire training command is wired for " |
| f"{workdir_rel} yet." |
| ) |
| return 0 if args.dry_run else SKIP_EXIT_CODE |
| _print_train_command(workdir, cmd, env) |
| if args.dry_run: |
| return _evaluate(model_name, args.dataset, args.gpu, dry_run=True) |
| code = 0 |
| if not args.eval_only: |
| log_dir = ROOT / "results" / model_name / dataset_cfg["name"] / "logs" |
| code = _run_logged(cmd, workdir, env, log_dir / "train_stdout.log", log_dir / "train_stderr.log") |
| if code == 0: |
| _canonicalize_checkpoints(model_name, dataset_cfg["name"]) |
| code = _evaluate(model_name, args.dataset, args.gpu, dry_run=False) |
| return code |
|
|
|
|
| def run_external_model(model_name: str, args: argparse.Namespace) -> int: |
| from utils.model_adapters import get_model_adapter |
| from utils.unified_trainer import train_with_adapter |
|
|
| adapter = get_model_adapter(model_name) |
| if adapter.supports_unified_training: |
| return train_with_adapter(model_name, args) |
| cfg = _dataset_ready(args.dataset, return_format="tuple", dry_run=args.dry_run) |
| apply_dataloader_cli_overrides(cfg, args) |
| print_dataloader_policy(cfg, torch.cuda.is_available()) |
| repo_rel, repo_url = NEW_REPOS[model_name] |
| repo = ROOT / repo_rel |
| if not repo.exists(): |
| print( |
| f"[SKIP] {model_name}/{args.dataset}: source repo is not present at {repo}. " |
| f"Expected source: {repo_url}." |
| ) |
| return 0 if args.dry_run else SKIP_EXIT_CODE |
| if model_name == "tinycd" and int(cfg.get("channels", 3)) > 3: |
| raise ValueError("TinyCD supports RGB input only. Set channels: 3 in the dataset config override.") |
| weight_path = None |
| if not args.dry_run: |
| weight_path = _ensure_pretrained(model_name) |
| print(f"[READY] {model_name} repo found at {repo}") |
| if weight_path: |
| print(f"[WEIGHTS] {weight_path}") |
| train_cmd, train_env, train_cwd = _external_train_command(model_name, cfg, args) |
| if not train_cmd: |
| print("[NOTE] Model-specific training command construction must be completed from the cloned repo config.") |
| print(f"[SKIP] {model_name}/{args.dataset}: repo-specific train command wiring is not implemented yet.") |
| return 0 if args.dry_run else SKIP_EXIT_CODE |
| if args.dry_run or args.eval_only: |
| _print_train_command(train_cwd, train_cmd, train_env) |
| return _evaluate(model_name, args.dataset, args.gpu, dry_run=args.dry_run) |
| _print_train_command(train_cwd, train_cmd, train_env) |
| log_dir = ROOT / "results" / model_name / cfg["name"] / "logs" |
| if model_name == "changemamba": |
| result_dir = ROOT / "results" / model_name / cfg["name"] |
| checkpoint_dir = result_dir / "checkpoints" |
| code = _run_streamed(train_cmd, train_cwd, train_env, log_dir / "train_wrapper.log") |
| if code != 0: |
| return FAILED_SUBPROCESS_EXIT_CODE |
| checkpoints = _checkpoint_files(checkpoint_dir) |
| if not checkpoints: |
| _print_checkpoint_failure(model_name, cfg["name"], train_cmd, checkpoint_dir, result_dir) |
| return FAILED_NO_CHECKPOINT_EXIT_CODE |
| _canonicalize_checkpoints(model_name, cfg["name"]) |
| code = _evaluate(model_name, args.dataset, args.gpu, dry_run=False) |
| return 0 if code == 0 else FAILED_EVAL_EXIT_CODE |
| code = _run_logged(train_cmd, train_cwd, train_env, log_dir / "train_stdout.log", log_dir / "train_stderr.log") |
| if code == 0 and model_name not in {"fc_ef", "fc_siam_conc", "fc_siam_diff"}: |
| _canonicalize_checkpoints(model_name, cfg["name"]) |
| code = _evaluate(model_name, args.dataset, args.gpu, dry_run=False) |
| return code |
|
|
|
|
| def run_smoke_test(model_name: str, args: argparse.Namespace) -> int: |
| from datasets.cd_dataset import CDDataset |
| from utils.config_loader import load_dataset_config, load_model_config |
| from utils.metrics import BinaryMetrics |
| from utils.model_adapters import get_model_adapter |
| from utils.profiling import count_flops, count_parameters |
| from utils.qualitative import safe_sample_id, save_binary_prediction, save_probability_map, select_or_load_manifest |
|
|
| cfg = load_dataset_config(args.dataset) |
| apply_dataloader_cli_overrides(cfg, args) |
| print_dataloader_policy(cfg, torch.cuda.is_available()) |
| root = Path(cfg["data_root"]) |
| print(f"[SMOKE] {model_name}/{args.dataset}") |
| print(f"[SMOKE] Dataset root exists: {root.is_dir()} -> {root}") |
| if not root.is_dir(): |
| return 1 |
| ds = CDDataset(root, "test", cfg=cfg, return_format="tuple") |
| a, b, mask, name = ds[0] |
| print(f"[SMOKE] Batch sample: name={name} A={tuple(a.shape)} B={tuple(b.shape)} mask={tuple(mask.shape)}") |
| adapter = get_model_adapter(model_name) |
| if adapter.supports_inprocess_eval: |
| gpu_resolution = resolve_gpu(args.gpu) |
| print_gpu_diagnostics(gpu_resolution) |
| device = torch.device(gpu_resolution.local_device) |
| model = adapter.build_model(load_model_config(model_name), cfg, device) |
| ckpt = ROOT / "results" / model_name / cfg["name"] / "checkpoints" / "best_model.pth" |
| if ckpt.exists(): |
| adapter.load_checkpoint(model, ckpt, device) |
| print(f"[SMOKE] Loaded checkpoint: {ckpt}") |
| else: |
| print(f"[SMOKE] No checkpoint found at {ckpt}; validating construction and forward pass only.") |
| batch = (a.unsqueeze(0), b.unsqueeze(0), mask.unsqueeze(0), [name]) |
| train_step = None |
| if adapter.supports_unified_training: |
| model.train() |
| optimizer = adapter.build_optimizer(model, load_model_config(model_name)) |
| optimizer.zero_grad(set_to_none=True) |
| raw_train = adapter.forward(model, batch, device) |
| loss_dict = adapter.compute_loss(raw_train, batch, load_model_config(model_name), cfg, device) |
| loss_dict["loss"].backward() |
| optimizer.step() |
| train_step = {key: float(value.detach().cpu().item()) for key, value in loss_dict.items()} |
| model.eval() |
| with torch.inference_mode(): |
| raw = adapter.forward(model, batch, device) |
| normalized = adapter.normalize_output(raw, batch, cfg) |
| metrics = BinaryMetrics(threshold=float(cfg.get("eval", {}).get("threshold", 0.5))) |
| metrics.update(normalized.metric_tensor, mask.unsqueeze(0)) |
| params = count_parameters(model) |
| flops = None |
| if adapter.supports_flops: |
| flops = count_flops(model, lambda: adapter.get_dummy_inputs(cfg, device), device) |
| manifest = select_or_load_manifest(cfg) |
| smoke_dir = ROOT / "results" / model_name / cfg["name"] / "smoke_test" |
| clean_id = safe_sample_id(str(name)) |
| save_binary_prediction(normalized.binary[0], smoke_dir / f"{clean_id}_pred.png") |
| if normalized.score is not None: |
| save_probability_map(normalized.score[0], smoke_dir / f"{clean_id}_prob.png") |
| print( |
| f"[SMOKE] PASS {model_name}/{cfg['name']} adapter={adapter.model_class_path} " |
| f"output={tuple(normalized.metric_tensor.shape)} params_m={params['params_m']:.3f} " |
| f"flops_g={None if flops is None else flops.get('flops_g')} " |
| f"manifest_samples={manifest.get('selected_count')} train_step={train_step} metrics={metrics.compute()}" |
| ) |
| return 0 |
| args.dry_run = True |
| if model_name in EXISTING_LEGACY: |
| workdir, cmd, env = _existing_train_command(model_name, cfg, args) |
| else: |
| cmd, env, workdir = _external_train_command(model_name, cfg, args) |
| if cmd: |
| print(_gpu_mapping_message(env)) |
| print(f"[SMOKE] Command construction OK: cwd={workdir} cmd={' '.join(cmd)}") |
| else: |
| print("[SMOKE] Command construction reports this adapter is not wired yet.") |
| print( |
| f"[SMOKE-UNAVAILABLE] {model_name}/{args.dataset}: {adapter.notes_or_failure_reason}" |
| ) |
| return 1 |
|
|
|
|
| def _external_train_command(model_name: str, dataset_cfg: dict, args: argparse.Namespace) -> tuple[list[str], dict, Path]: |
| from utils.legacy_config_writer import prepare_legacy_list_view |
|
|
| env, _ = subprocess_gpu_env(args.gpu) |
| child_gpu = "0" |
| model_cfg = _model_config(model_name) |
| out_dir = ROOT / "results" / model_name / dataset_cfg["name"] |
| out_dir.mkdir(parents=True, exist_ok=True) |
| if model_name == "dsamnet": |
| repo = ROOT / "model_repos" / "DSAMNet" |
| train_root = _split_dir(dataset_cfg, "train") |
| val_root = _split_dir(dataset_cfg, "val") |
| ckpt_dir = out_dir / "checkpoints" |
| ckpt_dir.mkdir(parents=True, exist_ok=True) |
| cmd = [ |
| sys.executable, |
| "train.py", |
| "--num_epochs", |
| str(int(model_cfg.get("num_epochs", 100))), |
| "--batchsize", |
| str(int(dataset_cfg.get("batch_size", 8))), |
| "--val_batchsize", |
| str(int(dataset_cfg.get("batch_size", 8))), |
| "--num_workers", |
| str(int(dataset_cfg.get("num_workers", 4))), |
| "--gpu_id", |
| child_gpu, |
| "--train1_dir", |
| str(train_root / dataset_cfg.get("image_a_folder", "A")), |
| "--train2_dir", |
| str(train_root / dataset_cfg.get("image_b_folder", "B")), |
| "--label_train", |
| str(train_root / dataset_cfg.get("mask_folder", "label")), |
| "--val1_dir", |
| str(val_root / dataset_cfg.get("image_a_folder", "A")), |
| "--val2_dir", |
| str(val_root / dataset_cfg.get("image_b_folder", "B")), |
| "--label_val", |
| str(val_root / dataset_cfg.get("mask_folder", "label")), |
| "--model_dir", |
| str(ckpt_dir) + "/", |
| "--sta_dir", |
| str(out_dir / "statistics.csv"), |
| ] |
| return cmd, env, repo |
| if model_name == "cgnet": |
| repo = ROOT / "model_repos" / "CGNet-CD" |
| if args.dry_run: |
| train_view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "train" |
| val_view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "val" |
| else: |
| train_view = _prepare_matched_split_view(dataset_cfg, "train") |
| val_view = _prepare_matched_split_view(dataset_cfg, "val") |
| cmd = [ |
| sys.executable, |
| "train_CGNet.py", |
| "--epoch", |
| str(int(model_cfg.get("num_epochs", 50))), |
| "--batchsize", |
| str(int(dataset_cfg.get("batch_size", 8))), |
| "--trainsize", |
| str(int(dataset_cfg.get("img_size", 256))), |
| "--gpu_id", |
| child_gpu, |
| "--data_name", |
| dataset_cfg["name"], |
| "--model_name", |
| "CGNet", |
| "--save_path", |
| str(out_dir) + "/", |
| "--train_root", |
| str(train_view) + "/", |
| "--val_root", |
| str(val_view) + "/", |
| ] |
| return cmd, env, repo |
| if model_name == "hanet": |
| repo = ROOT / "model_repos" / "HANet-CD" |
| view_root = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "hanet_matched" |
| if not args.dry_run: |
| for split in ("train", "val"): |
| split_view = _prepare_matched_split_view(dataset_cfg, split) |
| _safe_link_or_copy(split_view, view_root / split) |
| env["HANET_METADATA_JSON"] = str(_write_hanet_metadata(dataset_cfg, model_cfg, view_root)) |
| return [sys.executable, "trainHCX.py"], env, repo |
| if model_name == "tinycd": |
| repo = ROOT / "model_repos" / "Tiny_model_4_CD" |
| view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "tinycd" if args.dry_run else _prepare_tinycd_view(dataset_cfg) |
| cmd = [ |
| sys.executable, |
| "training.py", |
| "--datapath", |
| str(view), |
| "--log-path", |
| str(out_dir / "logs"), |
| "--gpu-id", |
| child_gpu, |
| "--batch-size", |
| str(int(args.batch_size or dataset_cfg.get("batch_size", 8))), |
| "--epochs", |
| str(int(args.epochs or model_cfg.get("num_epochs", 100))), |
| ] |
| return cmd, env, repo |
| if model_name == "changer": |
| repo = ROOT / "model_repos" / "open-cd" |
| view = ( |
| ROOT / "generated_dataset_views" / dataset_cfg["name"] / "opencd" |
| if args.dry_run |
| else _prepare_opencd_view(dataset_cfg) |
| ) |
| cfg_path = _write_opencd_changer_config(dataset_cfg, model_cfg, view) |
| return [sys.executable, "tools/train.py", str(cfg_path), "--work-dir", str(out_dir / "work_dir")], env, repo |
| if model_name == "changemamba": |
| repo = ROOT / "model_repos" / "ChangeMamba" |
| changedetection = repo / "changedetection" |
| view = ( |
| ROOT / "generated_dataset_views" / dataset_cfg["name"] / "changemamba" |
| if args.dry_run |
| else _prepare_changemamba_view(dataset_cfg) |
| ) |
| variant = str(model_cfg.get("vmamba_variant", "small")).lower() |
| model_type = str(model_cfg.get("changemamba_model_type", f"MambaBCD_{variant.capitalize()}")) |
| pretrained_name = { |
| "tiny": "vssmtiny_dp01_ckpt_epoch_292.pth", |
| "small": "vssmsmall_dp03_ckpt_epoch_238.pth", |
| "base": "vssmbase_dp06_ckpt_epoch_241.pth", |
| }.get(variant) |
| cfg_name = { |
| "tiny": "vssm_tiny_224_0229flex.yaml", |
| "small": "vssm_small_224.yaml", |
| "base": "vssm_base_224.yaml", |
| }.get(variant) |
| if pretrained_name is None or cfg_name is None: |
| raise ValueError(f"Unsupported ChangeMamba VMamba variant: {variant}") |
| pretrained = repo / "pretrained_weight" / pretrained_name |
| cfg_path = changedetection / "configs" / "vssm1" / cfg_name |
| max_iters, max_iters_source = _iteration_budget("changemamba", model_cfg, args, default=50000) |
| print(f"[TRAIN-CFG] smoke_test={bool(args.smoke_test)}") |
| print(f"[TRAIN-CFG] max_iters={max_iters}") |
| print(f"[TRAIN-CFG] max_iters_source={max_iters_source}") |
| cmd = [ |
| sys.executable, |
| "script/train_MambaBCD.py", |
| "--dataset", |
| _legacy_dataset_token(dataset_cfg), |
| "--batch_size", |
| str(int(dataset_cfg.get("batch_size", 8))), |
| "--num_workers", |
| str(int(dataset_cfg.get("num_workers", 4))), |
| "--crop_size", |
| str(int(dataset_cfg.get("img_size", model_cfg.get("img_size", 256)))), |
| "--max_iters", |
| str(max_iters), |
| "--model_type", |
| model_type, |
| "--model_param_path", |
| str(out_dir / "checkpoints"), |
| "--train_dataset_path", |
| str(view / "train"), |
| "--train_data_list_path", |
| str(view / "train_set.txt"), |
| "--test_dataset_path", |
| str(view / "test"), |
| "--test_data_list_path", |
| str(view / "test_set.txt"), |
| "--cfg", |
| str(cfg_path), |
| "--encoder_pretrained_path", |
| str(pretrained), |
| "--learning_rate", |
| str(float(model_cfg.get("lr", 6e-5))), |
| "--weight_decay", |
| str(float(model_cfg.get("weight_decay", 0.01))), |
| ] |
| return cmd, env, changedetection |
| if model_name in {"fc_ef", "fc_siam_conc", "fc_siam_diff"}: |
| cmd = [ |
| sys.executable, |
| str(ROOT / "train" / "fc_adapter.py"), |
| "--model", |
| model_name, |
| "--dataset", |
| dataset_cfg["name"], |
| "--gpu", |
| str(args.gpu), |
| ] |
| if args.epochs is not None: |
| cmd.extend(["--epochs", str(args.epochs)]) |
| if args.batch_size is not None: |
| cmd.extend(["--batch-size", str(args.batch_size)]) |
| if args.lr is not None: |
| cmd.extend(["--lr", str(args.lr)]) |
| if args.resume: |
| cmd.append("--resume") |
| if args.eval_only: |
| cmd.append("--eval-only") |
| if args.force: |
| cmd.append("--force") |
| if args.output_dir: |
| cmd.extend(["--output-dir", args.output_dir]) |
| return cmd, env, ROOT |
| if model_name == "elgcnet": |
| repo = ROOT / "model_repos" / "elgcnet" |
| view = ( |
| ROOT / "generated_dataset_views" / dataset_cfg["name"] / "legacy_list" |
| if args.dry_run |
| else prepare_legacy_list_view(dataset_cfg) |
| ) |
| env["CD_MODELS_DATASET_ROOT"] = str(view) |
| env["CD_MODELS_DSIFN_ROOT"] = str(view) |
| cmd = [ |
| sys.executable, |
| "main_cd.py", |
| "--gpu_ids", |
| child_gpu, |
| "--project_name", |
| f"{dataset_cfg['name']}-train-{model_name}", |
| "--data_name", |
| _legacy_dataset_token(dataset_cfg), |
| "--img_size", |
| str(int(dataset_cfg.get("img_size", model_cfg.get("img_size", 256)))), |
| "--batch_size", |
| str(int(dataset_cfg.get("batch_size", 8))), |
| "--num_workers", |
| str(int(dataset_cfg.get("num_workers", 4))), |
| "--max_epochs", |
| str(int(model_cfg.get("num_epochs", 200))), |
| "--optimizer", |
| str(model_cfg.get("optimizer", "adamw")), |
| "--lr", |
| str(float(model_cfg.get("lr", 0.00031))), |
| "--loss", |
| "ce", |
| "--net_G", |
| "ELGCNet", |
| ] |
| return cmd, env, repo |
| return [], env, ROOT |
|
|
|
|
| def main_for(model_name: str) -> int: |
| args = parse_args(model_name) |
| gpu_resolution = resolve_gpu(args.gpu) |
| print_gpu_diagnostics(gpu_resolution) |
| selected = args.model if model_name == "fc_variants" else model_name |
| if selected == "fc_variants": |
| print("[SKIP] fc_variants requires --model fc_ef, --model fc_siam_conc, or --model fc_siam_diff.") |
| return 0 if args.dry_run else SKIP_EXIT_CODE |
| if args.smoke_test: |
| return run_smoke_test(selected, args) |
| if selected in EXISTING_LEGACY: |
| return run_existing_model(selected, args) |
| return run_external_model(selected, args) |
|
|
