"""Train DINOv3 + DeepLabV3+ for seaweed segmentation. The script supports single GPU, torchrun DDP, AMP, gradient accumulation, checkpoint resume, and rank-0-only logging/checkpointing. """ from __future__ import annotations import argparse import json import os import random from datetime import datetime from pathlib import Path from typing import Any import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import torch import torch.distributed as dist import torch.nn as nn import torch.optim as optim from torch.cuda.amp import GradScaler, autocast from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler from tqdm import tqdm from dinov3_deeplabv3plus import DinoV3DeepLabV3Plus, SeaweedSegmentationLoss from seaweed_segmentation_dataset import SeaweedSegmentationDataset, get_train_transforms, get_val_transforms def load_json(path: str | Path) -> dict[str, Any]: with open(path, "r", encoding="utf-8") as f: return json.load(f) def is_dist() -> bool: return dist.is_available() and dist.is_initialized() def rank() -> int: return dist.get_rank() if is_dist() else 0 def world_size() -> int: return dist.get_world_size() if is_dist() else 1 def is_main_process() -> bool: return rank() == 0 def setup_distributed() -> tuple[torch.device, int]: local_rank = int(os.environ.get("LOCAL_RANK", "0")) distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1 if distributed: backend = "nccl" if os.name != "nt" else "gloo" rank_id = int(os.environ["RANK"]) size = int(os.environ["WORLD_SIZE"]) if os.environ.get("USE_LIBUV", "1") == "0": from datetime import timedelta store = dist.TCPStore( os.environ.get("MASTER_ADDR", "127.0.0.1"), int(os.environ.get("MASTER_PORT", "29500")), world_size=size, is_master=rank_id == 0, timeout=timedelta(seconds=180), use_libuv=False, ) dist.init_process_group(backend=backend, store=store, rank=rank_id, world_size=size) else: dist.init_process_group(backend=backend) if torch.cuda.is_available(): if distributed: torch.cuda.set_device(local_rank) device = torch.device(f"cuda:{local_rank}") else: device = torch.device("cpu") return device, local_rank def cleanup_distributed() -> None: if is_dist(): dist.barrier() dist.destroy_process_group() def seed_everything(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def resolve_path(path: str | None, roots: list[str]) -> str | None: if not path: return path candidate = Path(path) if candidate.exists(): return str(candidate) for root in roots: rooted = Path(root) / path if rooted.exists(): return str(rooted) return path def collate_skip_none(batch: list[dict[str, Any] | None]) -> dict[str, Any]: valid = [item for item in batch if item is not None] if not valid: raise RuntimeError("All samples in this batch failed to load.") images = torch.stack([item["image"] for item in valid], dim=0) masks = torch.stack([item["mask"] for item in valid], dim=0) filenames = [item["filename"] for item in valid] return {"image": images, "mask": masks, "filename": filenames} class SeaweedSegmentationTrainer: def __init__(self, config: dict[str, Any], device: torch.device, local_rank: int): self.config = config self.device = device self.local_rank = local_rank self.use_amp = bool(config.get("amp", True)) and device.type == "cuda" self.grad_accum_steps = int(config.get("grad_accum_steps", 1)) weight_roots = config.get("weight_search_roots", []) config["backbone_weights"] = resolve_path(config.get("backbone_weights"), weight_roots) seed_everything(int(config.get("seed", 42)) + rank()) model = DinoV3DeepLabV3Plus( num_classes=config["num_classes"], backbone_name=config["backbone_name"], pretrained=config["pretrained"], weights=config["backbone_weights"], use_4channel=config["use_4channel"], freeze_backbone=config.get("freeze_backbone", False), ).to(device) if is_dist(): self.model = DDP( model, device_ids=[local_rank] if device.type == "cuda" else None, find_unused_parameters=True, ) else: self.model = model self.criterion = SeaweedSegmentationLoss( num_classes=config["num_classes"], focal_alpha=config.get("focal_alpha", [1.0, 2.0]), focal_gamma=config.get("focal_gamma", 2), dice_weight=config.get("dice_weight", 0.5), focal_weight=config.get("focal_weight", 1.0), background_weight=config.get("background_weight", 1.0), foreground_weight=config.get("foreground_weight", 2.0), ).to(device) self.optimizer = self._create_optimizer() self.scheduler = self._create_scheduler() self.scaler = GradScaler(enabled=self.use_amp) self.train_loader, self.val_loader = self._create_data_loaders() self.train_history = { "train_loss": [], "train_focal": [], "train_dice": [], "val_loss": [], "val_focal": [], "val_dice": [], "val_iou": [], "val_accuracy": [], } self.best_val_loss = float("inf") self.best_val_iou = 0.0 self.start_epoch = 0 self.output_dir = Path(config["output_dir"]) if is_main_process(): self.output_dir.mkdir(parents=True, exist_ok=True) with open(self.output_dir / "config.json", "w", encoding="utf-8") as f: json.dump(config, f, indent=2, ensure_ascii=False) if config.get("resume"): self.load_checkpoint(config["resume"]) @property def raw_model(self) -> nn.Module: return self.model.module if isinstance(self.model, DDP) else self.model def _create_optimizer(self): raw_model = self.raw_model if self.config.get("freeze_backbone", False): params = raw_model.get_trainable_parameters() lr = self.config.get("decoder_lr", 1e-4) if is_main_process(): print(f"Frozen backbone mode. Trainable tensors: {len(params)}") return optim.AdamW(params, lr=lr, weight_decay=self.config.get("weight_decay", 1e-4)) backbone_params = list(raw_model.get_backbone_params()) decoder_params = list(raw_model.get_decoder_params()) return optim.AdamW( [ {"params": backbone_params, "lr": self.config.get("backbone_lr", 1e-5)}, {"params": decoder_params, "lr": self.config.get("decoder_lr", 1e-4)}, ], weight_decay=self.config.get("weight_decay", 1e-4), ) def _create_scheduler(self): scheduler_type = self.config.get("scheduler", "cosine") if scheduler_type == "cosine": return optim.lr_scheduler.CosineAnnealingWarmRestarts(self.optimizer, T_0=10, T_mult=2, eta_min=1e-7) if scheduler_type == "step": return optim.lr_scheduler.StepLR(self.optimizer, step_size=30, gamma=0.1) return None def _create_data_loaders(self): use_aug = self.config.get("use_data_augmentation", True) train_transform = ( get_train_transforms(self.config["image_size"], self.config["use_4channel"]) if use_aug else get_val_transforms(self.config["image_size"], self.config["use_4channel"]) ) train_dataset = SeaweedSegmentationDataset( image_dir=self.config["train_image_dir"], mask_dir=self.config["train_mask_dir"], transform=train_transform, target_size=self.config["image_size"], use_4channel=self.config["use_4channel"], ) val_dataset = SeaweedSegmentationDataset( image_dir=self.config["val_image_dir"], mask_dir=self.config["val_mask_dir"], transform=get_val_transforms(self.config["image_size"], self.config["use_4channel"]), target_size=self.config["image_size"], use_4channel=self.config["use_4channel"], ) if len(train_dataset) == 0 or len(val_dataset) == 0: raise RuntimeError( f"Empty dataset: train={len(train_dataset)}, val={len(val_dataset)}. " "Check image/mask directories in the config." ) train_sampler = DistributedSampler(train_dataset, shuffle=True) if is_dist() else None val_sampler = DistributedSampler(val_dataset, shuffle=False) if is_dist() else None num_workers = int(self.config.get("num_workers", 4)) loader_kwargs = { "batch_size": int(self.config["batch_size"]), "num_workers": num_workers, "pin_memory": self.device.type == "cuda", "collate_fn": collate_skip_none, "persistent_workers": num_workers > 0, } train_loader = DataLoader( train_dataset, shuffle=train_sampler is None, sampler=train_sampler, drop_last=True, **loader_kwargs, ) val_loader = DataLoader( val_dataset, shuffle=False, sampler=val_sampler, drop_last=False, **loader_kwargs, ) return train_loader, val_loader def calculate_metrics(self, pred, target): if isinstance(pred, dict): pred = pred["out"] pred_classes = torch.argmax(pred, dim=1) foreground_pred = pred_classes == 1 foreground_target = target == 1 intersection = (foreground_pred & foreground_target).sum().float() union = (foreground_pred | foreground_target).sum().float() iou = intersection / union.clamp_min(1) accuracy = (pred_classes == target).float().mean() return {"iou": iou.detach(), "accuracy": accuracy.detach()} def _reduce_scalar(self, value: torch.Tensor) -> float: value = value.detach().float() if is_dist(): dist.all_reduce(value, op=dist.ReduceOp.SUM) value /= world_size() return value.item() def run_epoch(self, epoch: int, train: bool): self.model.train(train) loader = self.train_loader if train else self.val_loader if train and isinstance(loader.sampler, DistributedSampler): loader.sampler.set_epoch(epoch) totals = {"total_loss": 0.0, "focal_loss": 0.0, "dice_loss": 0.0, "iou": 0.0, "accuracy": 0.0} steps = 0 desc = f"Epoch {epoch + 1}/{self.config['num_epochs']} - {'Train' if train else 'Val'}" iterator = tqdm(loader, desc=desc, disable=not is_main_process()) if train: self.optimizer.zero_grad(set_to_none=True) for step, batch in enumerate(iterator): max_batches = self.config.get("max_train_batches" if train else "max_val_batches") if max_batches is not None and step >= int(max_batches): break images = batch["image"].to(self.device, non_blocking=True) masks = batch["mask"].to(self.device, non_blocking=True) with torch.set_grad_enabled(train): with autocast(enabled=self.use_amp): outputs = self.model(images) losses = self.criterion(outputs, masks) loss = losses["total_loss"] / self.grad_accum_steps if train: self.scaler.scale(loss).backward() should_step = (step + 1) % self.grad_accum_steps == 0 or (step + 1) == len(loader) if should_step: if self.config.get("grad_clip"): self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.raw_model.parameters(), self.config["grad_clip"]) self.scaler.step(self.optimizer) self.scaler.update() self.optimizer.zero_grad(set_to_none=True) metrics = self.calculate_metrics(outputs, masks) totals["total_loss"] += self._reduce_scalar(losses["total_loss"]) totals["focal_loss"] += self._reduce_scalar(losses["focal_loss"]) totals["dice_loss"] += self._reduce_scalar(losses["dice_loss"]) totals["iou"] += self._reduce_scalar(metrics["iou"]) totals["accuracy"] += self._reduce_scalar(metrics["accuracy"]) steps += 1 if is_main_process(): iterator.set_postfix( loss=f"{totals['total_loss'] / steps:.4f}", iou=f"{totals['iou'] / steps:.4f}", acc=f"{totals['accuracy'] / steps:.4f}", ) return {k: v / max(steps, 1) for k, v in totals.items()} def save_checkpoint(self, epoch: int, is_best: bool) -> None: if not is_main_process(): return checkpoint = { "epoch": epoch, "model_state_dict": self.raw_model.state_dict(), "optimizer_state_dict": self.optimizer.state_dict(), "scheduler_state_dict": self.scheduler.state_dict() if self.scheduler else None, "scaler_state_dict": self.scaler.state_dict(), "train_history": self.train_history, "config": self.config, "best_val_loss": self.best_val_loss, "best_val_iou": self.best_val_iou, } torch.save(checkpoint, self.output_dir / "latest_checkpoint.pth") if is_best: torch.save(checkpoint, self.output_dir / "best_checkpoint.pth") print(f"Saved best checkpoint: {self.output_dir / 'best_checkpoint.pth'}") def load_checkpoint(self, path: str) -> None: checkpoint = torch.load(path, map_location=self.device) self.raw_model.load_state_dict(checkpoint["model_state_dict"], strict=True) self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) if self.scheduler and checkpoint.get("scheduler_state_dict"): self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) if checkpoint.get("scaler_state_dict"): self.scaler.load_state_dict(checkpoint["scaler_state_dict"]) self.train_history = checkpoint.get("train_history", self.train_history) self.best_val_loss = checkpoint.get("best_val_loss", self.best_val_loss) self.best_val_iou = checkpoint.get("best_val_iou", self.best_val_iou) self.start_epoch = int(checkpoint.get("epoch", -1)) + 1 if is_main_process(): print(f"Resumed from {path} at epoch {self.start_epoch}") def plot_training_history(self) -> None: if not is_main_process() or not self.train_history["train_loss"]: return epochs = range(1, len(self.train_history["train_loss"]) + 1) fig, axes = plt.subplots(2, 2, figsize=(15, 10)) axes[0, 0].plot(epochs, self.train_history["train_loss"], label="Train") axes[0, 0].plot(epochs, self.train_history["val_loss"], label="Val") axes[0, 0].set_title("Loss") axes[0, 0].legend() axes[0, 1].plot(epochs, self.train_history["train_focal"], label="Train") axes[0, 1].plot(epochs, self.train_history["val_focal"], label="Val") axes[0, 1].set_title("Focal Loss") axes[0, 1].legend() axes[1, 0].plot(epochs, self.train_history["train_dice"], label="Train") axes[1, 0].plot(epochs, self.train_history["val_dice"], label="Val") axes[1, 0].set_title("Dice Loss") axes[1, 0].legend() axes[1, 1].plot(epochs, self.train_history["val_iou"], label="Val IoU") axes[1, 1].plot(epochs, self.train_history["val_accuracy"], label="Val Acc") axes[1, 1].set_title("Metrics") axes[1, 1].legend() for ax in axes.ravel(): ax.grid(True) plt.tight_layout() plt.savefig(self.output_dir / "training_history.png", dpi=200, bbox_inches="tight") plt.close(fig) def train(self): if is_main_process(): effective_batch = self.config["batch_size"] * world_size() * self.grad_accum_steps print(f"Device: {self.device}; world_size={world_size()}; AMP={self.use_amp}") print(f"Per-GPU batch: {self.config['batch_size']}; effective batch: {effective_batch}") print(f"Train samples: {len(self.train_loader.dataset)}; val samples: {len(self.val_loader.dataset)}") print(f"Parameters: {sum(p.numel() for p in self.raw_model.parameters()):,}") for epoch in range(self.start_epoch, int(self.config["num_epochs"])): train_stats = self.run_epoch(epoch, train=True) val_stats = self.run_epoch(epoch, train=False) if self.scheduler: self.scheduler.step() self.train_history["train_loss"].append(train_stats["total_loss"]) self.train_history["train_focal"].append(train_stats["focal_loss"]) self.train_history["train_dice"].append(train_stats["dice_loss"]) self.train_history["val_loss"].append(val_stats["total_loss"]) self.train_history["val_focal"].append(val_stats["focal_loss"]) self.train_history["val_dice"].append(val_stats["dice_loss"]) self.train_history["val_iou"].append(val_stats["iou"]) self.train_history["val_accuracy"].append(val_stats["accuracy"]) is_best = val_stats["total_loss"] < self.best_val_loss or val_stats["iou"] > self.best_val_iou self.best_val_loss = min(self.best_val_loss, val_stats["total_loss"]) self.best_val_iou = max(self.best_val_iou, val_stats["iou"]) if is_main_process(): print( f"Epoch {epoch + 1}: train_loss={train_stats['total_loss']:.4f}, " f"val_loss={val_stats['total_loss']:.4f}, val_iou={val_stats['iou']:.4f}" ) self.save_checkpoint(epoch, is_best) if (epoch + 1) % int(self.config.get("plot_interval", 10)) == 0: self.plot_training_history() self.plot_training_history() if is_main_process(): print(f"Done. best_val_loss={self.best_val_loss:.4f}, best_val_iou={self.best_val_iou:.4f}") def default_config() -> dict[str, Any]: return { "train_image_dir": "data/train/images", "train_mask_dir": "data/train/masks", "val_image_dir": "data/val/images", "val_mask_dir": "data/val/masks", "num_classes": 2, "backbone_name": "dinov3_vitl16", "pretrained": True, "backbone_weights": "dinov3_vitl16_pretrain_sat493m-eadcf0ff.pth", "weight_search_roots": ["D:/DINOv3_pretrained_weights", "D:/DINOv3预训练权重", "dinov3-main"], "use_4channel": True, "image_size": 256, "batch_size": 8, "num_epochs": 100, "backbone_lr": 1e-5, "decoder_lr": 1e-4, "weight_decay": 1e-4, "scheduler": "cosine", "grad_clip": 1.0, "grad_accum_steps": 1, "amp": True, "focal_alpha": [1.0, 3.0], "focal_gamma": 2, "dice_weight": 0.5, "focal_weight": 1.0, "background_weight": 1.0, "foreground_weight": 3.0, "num_workers": 4, "output_dir": f"outputs/seaweed_segmentation_{datetime.now().strftime('%Y%m%d_%H%M%S')}", "plot_interval": 10, "seed": 42, "freeze_backbone": False, "use_data_augmentation": True, } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Train seaweed segmentation") parser.add_argument("--config", type=str, default=None, help="JSON config path") parser.add_argument("--resume", type=str, default=None, help="Checkpoint to resume") parser.add_argument("--batch-size", type=int, default=None, help="Override per-GPU batch size") parser.add_argument("--epochs", type=int, default=None, help="Override epoch count") parser.add_argument("--output-dir", type=str, default=None, help="Override output directory") parser.add_argument("--weights", type=str, default=None, help="Override DINOv3 weights path") parser.add_argument("--max-train-batches", type=int, default=None, help="Debug limit for train batches per epoch") parser.add_argument("--max-val-batches", type=int, default=None, help="Debug limit for val batches per epoch") return parser.parse_args() def main() -> None: args = parse_args() config = default_config() if args.config: config.update(load_json(args.config)) if args.resume: config["resume"] = args.resume if args.batch_size: config["batch_size"] = args.batch_size if args.epochs: config["num_epochs"] = args.epochs if args.output_dir: config["output_dir"] = args.output_dir if args.weights: config["backbone_weights"] = args.weights if args.max_train_batches is not None: config["max_train_batches"] = args.max_train_batches if args.max_val_batches is not None: config["max_val_batches"] = args.max_val_batches device, local_rank = setup_distributed() try: trainer = SeaweedSegmentationTrainer(config, device, local_rank) trainer.train() finally: cleanup_distributed() if __name__ == "__main__": main()