Spaces:
Sleeping
Sleeping
| import argparse | |
| import os | |
| from typing import List, Tuple | |
| import torch | |
| # Choices for ``--dino_variant``. Centralised so train and test cannot drift. | |
| DINO_CHOICES = [ | |
| "dinov2-base", | |
| "dinov2-large", | |
| "dinov2-giant", | |
| "dinov3-vitb16-pretrain-lvd1689m", | |
| ] | |
| # Choices for ``--optim``. | |
| OPTIM_CHOICES = ["adam", "sgd"] | |
| def build_train_parser() -> argparse.ArgumentParser: | |
| """Construct the training CLI parser. | |
| Every option used anywhere in the training pipeline is registered here | |
| with a sensible default; downstream modules access them directly via | |
| attribute lookup. | |
| """ | |
| parser = argparse.ArgumentParser( | |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter, | |
| description="PGC (Peak-Guided Calibration) training CLI", | |
| ) | |
| # ----- Identity / I/O ------------------------------------------------- | |
| parser.add_argument( | |
| "--name", | |
| type=str, | |
| default="pgc_train", | |
| help="Experiment name (also used as log filename stem).", | |
| ) | |
| parser.add_argument( | |
| "--checkpoints_dir", | |
| type=str, | |
| default="checkpoints", | |
| help="Directory under which logs and checkpoints are written.", | |
| ) | |
| parser.add_argument( | |
| "--devices", | |
| type=str, | |
| default="3", | |
| help=( | |
| "CUDA devices, e.g. '0' or '0,1'; '-1' for CPU. " | |
| "RTX 4090 #3 is the test box." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--num_threads", | |
| type=int, | |
| default=8, | |
| help="Number of DataLoader worker processes for training.", | |
| ) | |
| parser.add_argument( | |
| "--batch_size", | |
| type=int, | |
| default=32, | |
| help="Training batch size.", | |
| ) | |
| # ----- DINO backbone & LoRA ------------------------------------------- | |
| parser.add_argument( | |
| "--dino_variant", | |
| type=str, | |
| default="dinov2-large", | |
| choices=DINO_CHOICES, | |
| help="DINO backbone variant.", | |
| ) | |
| parser.add_argument( | |
| "--dino_pretrained_root", | |
| type=str, | |
| default=None, | |
| help=( | |
| "Root directory containing local DINO checkpoints. " | |
| "When set, the loader expects subfolders named after the variant. " | |
| "Required for training; entry-point validates explicitly." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--lora_rank", | |
| type=int, | |
| default=8, | |
| help="LoRA rank for adapter layers.", | |
| ) | |
| parser.add_argument( | |
| "--lora_alpha", | |
| type=float, | |
| default=1.0, | |
| help="LoRA scaling factor (alpha / rank determines effective scale).", | |
| ) | |
| parser.add_argument( | |
| "--lora_dropout", | |
| type=float, | |
| default=0.1, | |
| help="Dropout probability inside LoRA adapter layers.", | |
| ) | |
| # ----- Training data --------------------------------------------------- | |
| # ``--real_image_dir`` and ``--fake_image_dir`` both accept *one or more* | |
| # root directories (``nargs='+'``). ``RealFakeDataset`` scans every | |
| # supplied directory, concatenates the samples and performs a global | |
| # ``random.shuffle`` so that every mini-batch can freely mix sources. | |
| parser.add_argument( | |
| "--real_image_dir", | |
| type=str, | |
| nargs="+", | |
| default=None, | |
| help=( | |
| "One or more root directories of real training images. " | |
| "Example: --real_image_dir /path/A /path/B. All directories are " | |
| "scanned recursively and the resulting samples are merged + " | |
| "shuffled before training." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--fake_image_dir", | |
| type=str, | |
| nargs="+", | |
| default=None, | |
| help=( | |
| "One or more root directories of fake / generated training " | |
| "images. Example: --fake_image_dir /path/A /path/B. All " | |
| "directories are scanned recursively and the resulting samples " | |
| "are merged + shuffled before training." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--cropSize", | |
| type=int, | |
| default=224, | |
| help="Pad-then-crop input resolution (square).", | |
| ) | |
| # ----- Mid-training evaluation ---------------------------------------- | |
| parser.add_argument( | |
| "--test_root", | |
| type=str, | |
| default=None, | |
| help=( | |
| "Optional root directory of the evaluation dataset " | |
| "(UniversalFakeDetect layout: ``<root>/<subset>/{0_real,1_fake}``). " | |
| "If supplied the trainer evaluates on it periodically." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--eval_batch_size", | |
| type=int, | |
| default=32, | |
| help="Batch size used during mid-training evaluation.", | |
| ) | |
| parser.add_argument( | |
| "--eval_num_threads", | |
| type=int, | |
| default=8, | |
| help="DataLoader workers for mid-training evaluation.", | |
| ) | |
| # ----- Optimization ---------------------------------------------------- | |
| parser.add_argument( | |
| "--niter", | |
| type=int, | |
| default=100, | |
| help="Total number of training epochs.", | |
| ) | |
| parser.add_argument( | |
| "--lr", | |
| type=float, | |
| default=5e-5, | |
| help="Initial learning rate.", | |
| ) | |
| parser.add_argument( | |
| "--optim", | |
| type=str, | |
| default="adam", | |
| choices=OPTIM_CHOICES, | |
| help="Optimizer family.", | |
| ) | |
| parser.add_argument( | |
| "--accumulation_steps", | |
| type=int, | |
| default=4, | |
| help="Gradient accumulation steps.", | |
| ) | |
| parser.add_argument( | |
| "--weight_decay", | |
| type=float, | |
| default=0.05, | |
| help="L2 regularization strength.", | |
| ) | |
| parser.add_argument( | |
| "--label_smoothing", | |
| type=float, | |
| default=0.1, | |
| help="Label-smoothing epsilon for binary cross-entropy.", | |
| ) | |
| # ----- PGCM (Peak-Guided Calibration Module) hyper-parameters --------- | |
| parser.add_argument( | |
| "--tau_rgb", | |
| type=float, | |
| default=0.5, | |
| help="Temperature for the RGB-stream peak aggregation in PGCM.", | |
| ) | |
| parser.add_argument( | |
| "--tau_res", | |
| type=float, | |
| default=0.5, | |
| help="Temperature for the residual-stream peak aggregation in PGCM.", | |
| ) | |
| # ----- Stability / reproducibility ------------------------------------ | |
| parser.add_argument( | |
| "--max_grad_norm", | |
| type=float, | |
| default=1.0, | |
| help="Maximum L2 norm for gradient clipping (0 to disable).", | |
| ) | |
| parser.add_argument( | |
| "--seed", | |
| type=int, | |
| default=42, | |
| help="Random seed for Python, NumPy and PyTorch.", | |
| ) | |
| # ----- Bookkeeping ----------------------------------------------------- | |
| parser.add_argument( | |
| "--eval_every_steps", | |
| type=int, | |
| default=200, | |
| help=( | |
| "Run mid-training evaluation every N optimization steps. " | |
| "A checkpoint named " | |
| "``step{N}_acc{int}_ap{int}.pth`` is also saved right after " | |
| "each such evaluation (in addition to the best-AP ``best.pth``)." | |
| ), | |
| ) | |
| # Internal flag used by ``finalize_opt`` to know which parser produced | |
| # the namespace. Not exposed as a CLI option (no ``add_argument``). | |
| parser.set_defaults(isTrain=True) | |
| return parser | |
| def build_test_parser() -> argparse.ArgumentParser: | |
| """Construct the inference / evaluation CLI parser.""" | |
| parser = argparse.ArgumentParser( | |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter, | |
| description="PGC (Peak-Guided Calibration) test CLI", | |
| ) | |
| parser.add_argument( | |
| "--checkpoint", | |
| type=str, | |
| required=True, | |
| help="Path to the trained model checkpoint (.pth).", | |
| ) | |
| parser.add_argument( | |
| "--test_root", | |
| type=str, | |
| required=True, | |
| help="Root directory of the evaluation dataset.", | |
| ) | |
| parser.add_argument( | |
| "--checkpoints_dir", | |
| type=str, | |
| default="checkpoints", | |
| help="Directory under which the test log file is written.", | |
| ) | |
| parser.add_argument( | |
| "--name", | |
| type=str, | |
| default="pgc_test", | |
| help="Experiment name used to derive the test log filename.", | |
| ) | |
| parser.add_argument( | |
| "--batch_size", | |
| type=int, | |
| default=32, | |
| help="Evaluation batch size.", | |
| ) | |
| parser.add_argument( | |
| "--num_workers", | |
| type=int, | |
| default=8, | |
| help="DataLoader worker processes for evaluation.", | |
| ) | |
| parser.add_argument( | |
| "--devices", | |
| type=str, | |
| default="3", | |
| help=( | |
| "CUDA devices, e.g. '0' or '0,1'; '-1' for CPU. " | |
| "RTX 4090 #3 is the test box." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--cropSize", | |
| type=int, | |
| default=224, | |
| help="Pad-then-crop input resolution (must match training).", | |
| ) | |
| parser.add_argument( | |
| "--dino_variant", | |
| type=str, | |
| default="dinov2-large", | |
| choices=DINO_CHOICES, | |
| help="DINO backbone variant (must match training).", | |
| ) | |
| parser.add_argument( | |
| "--lora_rank", | |
| type=int, | |
| default=8, | |
| help="LoRA rank (must match training).", | |
| ) | |
| parser.add_argument( | |
| "--lora_alpha", | |
| type=float, | |
| default=1.0, | |
| help="LoRA alpha (must match training).", | |
| ) | |
| parser.add_argument( | |
| "--lora_dropout", | |
| type=float, | |
| default=0.1, | |
| help="LoRA dropout (must match training).", | |
| ) | |
| parser.add_argument( | |
| "--dino_pretrained_root", | |
| type=str, | |
| default=None, | |
| help="Root directory containing local DINO checkpoints.", | |
| ) | |
| parser.add_argument( | |
| "--tau_rgb", | |
| type=float, | |
| default=0.5, | |
| help="PGCM RGB-stream temperature (must match training).", | |
| ) | |
| parser.add_argument( | |
| "--tau_res", | |
| type=float, | |
| default=0.5, | |
| help="PGCM residual-stream temperature (must match training).", | |
| ) | |
| parser.set_defaults(isTrain=False) | |
| return parser | |
| def parse_devices(devices_str: str) -> Tuple[List[int], torch.device]: | |
| str_ids = [s for s in devices_str.split(",") if s.strip() != ""] | |
| parsed_ids = [int(x) for x in str_ids] | |
| valid_ids = [i for i in parsed_ids if i >= 0] | |
| if torch.cuda.is_available() and len(valid_ids) > 0: | |
| return valid_ids, torch.device(f"cuda:{valid_ids[0]}") | |
| return [], torch.device("cpu") | |
| def cli_message_for_opt( | |
| parser: argparse.ArgumentParser, opt: argparse.Namespace | |
| ) -> str: | |
| message = "----------------- Options ---------------\n" | |
| for k, v in sorted(vars(opt).items()): | |
| if k.startswith("_"): | |
| continue | |
| default = parser.get_default(k) | |
| comment = f"\t[default: {default}]" if v != default else "" | |
| message += "{:>25}: {:<30}{}\n".format(str(k), str(v), comment) | |
| message += "----------------- End -------------------" | |
| return message | |
| def finalize_opt(opt: argparse.Namespace) -> str: | |
| parser = build_train_parser() if opt.isTrain else build_test_parser() | |
| message = cli_message_for_opt(parser, opt) | |
| expr_dir = os.path.join(opt.checkpoints_dir, opt.name) | |
| os.makedirs(expr_dir, exist_ok=True) | |
| with open(os.path.join(expr_dir, "opt.txt"), "wt") as opt_file: | |
| opt_file.write(message + "\n") | |
| devices_list, device = parse_devices(opt.devices) | |
| opt.devices = devices_list | |
| opt.device = device | |
| if torch.cuda.is_available() and len(devices_list) > 0: | |
| torch.cuda.set_device(devices_list[0]) | |
| return message | |