# Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the Apache License, Version 2.0 # found in the LICENSE file in the root directory of this source tree. import argparse import logging import math import os import random from functools import partial import itertools from io import BytesIO from pathlib import Path import gc import contextlib import glob from fvcore.common.checkpoint import PeriodicCheckpointer import torch from dinov2.data import collate_data_and_cast, DataAugmentationDINO, MaskingGenerator import dinov2.distributed as distributed from dinov2.fsdp import FSDPCheckpointer from dinov2.logging import MetricLogger from dinov2.utils.config import setup from dinov2.utils.utils import CosineScheduler from dinov2.train.ssl_meta_arch import SSLMetaArch from dinov2.models import build_model_from_cfg from datasets import IterableDatasetDict, load_dataset, DownloadConfig from PIL import Image, ImageOps import h5py import numpy as np import torch.utils.data from torchvision import transforms from torchvision.datasets import folder from tqdm import tqdm import pyarrow import pyarrow.dataset import torch.distributed as dist import torch.nn.functional as F import pyarrow import pyarrow.dataset import torch.distributed as dist from torch.distributed.fsdp import ( FullyShardedDataParallel as FSDP, StateDictType, FullStateDictConfig, ) torch.backends.cuda.matmul.allow_tf32 = True # PyTorch 1.12 sets this to False by default logger = logging.getLogger("dinov2") import wandb RUN_SCRIPT = os.environ.get("DINOV2_RUN_SCRIPT", "") AUGMENTATION_FILE = Path(__file__).resolve().parents[1] / "data" / "augmentations.py" VISION_TRANSFORMER_FILE = Path(__file__).resolve().parents[1] / "models" / "vision_transformer.py" SSL_META_ARCH = Path(__file__).resolve().parents[1] / "train" / "ssl_meta_arch.py" def _build_streaming_dataset( dataset_path: str, *, shuffle_buffer: int, base_seed: int = 0, fragment_prefetch_limit: int = 4, fragment_range_size: int = 32 << 20, epoch: int = 0, ): # Get current rank/size at call time (safe under elastic restarts) world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 global_rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 fragment_scan_options = pyarrow.dataset.ParquetFragmentScanOptions( cache_options=pyarrow.CacheOptions( prefetch_limit=fragment_prefetch_limit, range_size_limit=fragment_range_size, ), ) # 로컬 parquet glob(Test B: 다운로드된 TCGA-12K) vs HF repo id 자동 분기 if "*" in dataset_path or dataset_path.startswith("/"): ds = load_dataset("parquet", data_files=dataset_path, streaming=True, split="train") else: ds = load_dataset( dataset_path, streaming=True, fragment_scan_options=fragment_scan_options, )["train"] # 1) shard first to avoid cross-rank duplication and wasted I/O if world_size > 1: ds = ds.shard(num_shards=world_size, index=global_rank) # 2) then shuffle; vary by epoch and rank seed = base_seed + epoch * 1_000_000 + global_rank * 10000 ds = ds.shuffle(buffer_size=shuffle_buffer, seed=seed) return ds def get_args_parser(add_help: bool = True): parser = argparse.ArgumentParser("DINOv2 training", add_help=add_help) parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file") parser.add_argument( "--no-resume", action="store_true", help="Whether to not attempt to resume from the checkpoint directory. ", ) parser.add_argument("--eval-only", action="store_true", help="perform evaluation only") parser.add_argument("--eval", type=str, default="", help="Eval type to perform") parser.add_argument( "opts", help=""" Modify config options at the end of the command. For Yacs configs, use space-separated "PATH.KEY VALUE" pairs. For python-based LazyConfig, use "path.key=value". """.strip(), default=None, nargs=argparse.REMAINDER, ) parser.add_argument( "--output-dir", "--output_dir", default="", type=str, help="Output directory to save logs and checkpoints", ) return parser def build_optimizer(cfg, params_groups): return torch.optim.AdamW(params_groups, betas=(cfg.optim.adamw_beta1, cfg.optim.adamw_beta2)) def build_schedulers(cfg): OFFICIAL_EPOCH_LENGTH = cfg.train.OFFICIAL_EPOCH_LENGTH lr = dict( base_value=cfg.optim["lr"], final_value=cfg.optim["min_lr"], total_iters=cfg.optim["epochs"] * OFFICIAL_EPOCH_LENGTH, warmup_iters=cfg.optim["warmup_epochs"] * OFFICIAL_EPOCH_LENGTH, start_warmup_value=0, ) wd = dict( base_value=cfg.optim["weight_decay"], final_value=cfg.optim["weight_decay_end"], total_iters=cfg.optim["epochs"] * OFFICIAL_EPOCH_LENGTH, ) momentum = dict( base_value=cfg.teacher["momentum_teacher"], final_value=cfg.teacher["final_momentum_teacher"], total_iters=cfg.optim["epochs"] * OFFICIAL_EPOCH_LENGTH, ) teacher_temp = dict( base_value=cfg.teacher["teacher_temp"], final_value=cfg.teacher["teacher_temp"], total_iters=cfg.teacher["warmup_teacher_temp_epochs"] * OFFICIAL_EPOCH_LENGTH, warmup_iters=cfg.teacher["warmup_teacher_temp_epochs"] * OFFICIAL_EPOCH_LENGTH, start_warmup_value=cfg.teacher["warmup_teacher_temp"], ) lr_schedule = CosineScheduler(**lr) wd_schedule = CosineScheduler(**wd) momentum_schedule = CosineScheduler(**momentum) teacher_temp_schedule = CosineScheduler(**teacher_temp) last_layer_lr_schedule = CosineScheduler(**lr) last_layer_lr_schedule.schedule[ : cfg.optim["freeze_last_layer_epochs"] * OFFICIAL_EPOCH_LENGTH ] = 0 # mimicking the original schedules # ★ Gram anchoring 가중치 스케줄: it_first_update까지 0 → ramp 구간 선형 상승 → 이후 상수. class _GramWeightSchedule: def __init__(self, first_update, weight, ramp=5000): self.first_update = int(first_update); self.weight = float(weight); self.ramp = max(int(ramp), 1) def __getitem__(self, it): if not self.first_update or it < self.first_update: return 0.0 return self.weight * min(1.0, (it - self.first_update) / self.ramp) if getattr(cfg, "gram", None) and cfg.gram.use_loss: gram_weight_schedule = _GramWeightSchedule( cfg.gram.it_first_update, cfg.gram.loss_weight, cfg.gram.get("ramp_iters", 5000) ) else: gram_weight_schedule = _GramWeightSchedule(0, 0.0) # 항상 0 logger.info("Schedulers ready.") return ( lr_schedule, wd_schedule, momentum_schedule, teacher_temp_schedule, last_layer_lr_schedule, gram_weight_schedule, ) def apply_optim_scheduler(optimizer, lr, wd, last_layer_lr): for param_group in optimizer.param_groups: is_last_layer = param_group["is_last_layer"] lr_multiplier = param_group["lr_multiplier"] wd_multiplier = param_group["wd_multiplier"] param_group["weight_decay"] = wd * wd_multiplier param_group["lr"] = (last_layer_lr if is_last_layer else lr) * lr_multiplier _FULL_STATE_DICT_CFG = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) _HUB_ARCH_NAMES = { "vit_small": "vits", "vit_base": "vitb", "vit_large": "vitl", "vit_giant2": "vitg", } def _resolve_torchhub_name(cfg): hub_arch = _HUB_ARCH_NAMES.get(cfg.student.arch) if hub_arch is None: raise AssertionError(f"Unsupported student arch for pretrained load: {cfg.student.arch}") if cfg.student.patch_size != 14: raise AssertionError("Pretrained torch.hub weights are only defined for patch_size=14") if cfg.student.num_register_tokens not in (0, 4): raise AssertionError("Pretrained weights only support 0 or 4 register tokens") reg_suffix = "_reg" if cfg.student.num_register_tokens else "" return f"dinov2_{hub_arch}{cfg.student.patch_size}{reg_suffix}" def _iter_vit_blocks(backbone): if backbone.chunked_blocks: for chunk in backbone.blocks: for blk in chunk: if not isinstance(blk, torch.nn.Identity): yield blk else: for blk in backbone.blocks: yield blk def _load_pretrained_backbone(cfg, model): def _mlp_kind(block): if hasattr(block.mlp, "fc1"): return "mlp" if hasattr(block.mlp, "w12"): return "swiglu" raise AssertionError("Unsupported FFN block type") hub_name = _resolve_torchhub_name(cfg) logger.info("Building pretrained DINOv2 backbone (%s) via dinov2.hub.backbones", hub_name) from dinov2.hub import backbones as _dinov2_hub_backbones model_pretrained = getattr(_dinov2_hub_backbones, hub_name)(pretrained=True) device = next(model.parameters()).device model_pretrained = model_pretrained.to(device) student_backbone = model.student.backbone teacher_backbone = model.teacher.backbone with torch.no_grad(): if student_backbone.embed_dim != model_pretrained.embed_dim: raise AssertionError("Pretrained embed_dim mismatch") if student_backbone.n_blocks != model_pretrained.n_blocks: raise AssertionError("Pretrained depth mismatch") if student_backbone.num_register_tokens != model_pretrained.num_register_tokens: raise AssertionError("Pretrained register token count mismatch") student_backbone.patch_embed.proj.weight.copy_(model_pretrained.patch_embed.proj.weight) student_backbone.patch_embed.proj.bias.copy_(model_pretrained.patch_embed.proj.bias) student_backbone.cls_token.copy_(model_pretrained.cls_token) student_backbone.mask_token.copy_(model_pretrained.mask_token) if student_backbone.num_register_tokens: student_backbone.register_tokens.copy_(model_pretrained.register_tokens) pos_embed_pretrained = model_pretrained.pos_embed.detach() n_extra_tokens = 1 cls_pos_embed = pos_embed_pretrained[:, :n_extra_tokens] patch_pos_embed = pos_embed_pretrained[:, n_extra_tokens:] orig_size = int(patch_pos_embed.shape[1] ** 0.5) patch_pos_embed = patch_pos_embed.reshape(1, orig_size, orig_size, -1).permute(0, 3, 1, 2) target_h, target_w = student_backbone.patch_embed.patches_resolution resized_patch_pos_embed = F.interpolate( patch_pos_embed, size=(target_h, target_w), mode="bicubic", align_corners=False, antialias=model_pretrained.interpolate_antialias, ) resized_patch_pos_embed = resized_patch_pos_embed.permute(0, 2, 3, 1).reshape( 1, target_h * target_w, -1 ) new_pos_embed = torch.cat((cls_pos_embed, resized_patch_pos_embed), dim=1) student_backbone.pos_embed.copy_(new_pos_embed) teacher_backbone.pos_embed.copy_(new_pos_embed) student_blocks = list(_iter_vit_blocks(student_backbone)) pretrained_blocks = list(_iter_vit_blocks(model_pretrained)) if len(student_blocks) != len(pretrained_blocks): raise AssertionError("Pretrained block count mismatch") if student_blocks and _mlp_kind(student_blocks[0]) != _mlp_kind(pretrained_blocks[0]): raise AssertionError( f"FFN mismatch: cfg.student.ffn_layer builds {_mlp_kind(student_blocks[0])}, " f"but torch.hub {hub_name} uses {_mlp_kind(pretrained_blocks[0])}" ) for dst, src in zip(student_blocks, pretrained_blocks): dst.load_state_dict(src.state_dict(), strict=True) student_backbone.norm.weight.copy_(model_pretrained.norm.weight) student_backbone.norm.bias.copy_(model_pretrained.norm.bias) def _freeze_student_backbone_except_last_n(cfg, model): n_unfrozen = cfg.train.unfreeze_last_n_blocks student_backbone = model.student.backbone blocks = list(_iter_vit_blocks(student_backbone)) total_blocks = len(blocks) if n_unfrozen < 1 or n_unfrozen > total_blocks: n_unfrozen = total_blocks print("cfg.train.unfreeze_last_n_blocks not set; setting to total blocks which is {total_blocks}") freeze_until = total_blocks - n_unfrozen if freeze_until == 0: return for p in student_backbone.parameters(): p.requires_grad = False for blk in blocks[freeze_until:]: for p in blk.parameters(): p.requires_grad = True for p in student_backbone.norm.parameters(): p.requires_grad = True logger.info("Froze %d/%d student backbone blocks (trainable blocks: %d)", freeze_until, total_blocks, n_unfrozen) def do_test(cfg, model, iteration): # save teacher checkpoint (used for eval only) # All ranks participate in FSDP state_dict() even with rank0_only=True is_main = distributed.is_main_process() iterstring = str(iteration) eval_dir = os.path.join(cfg.train.output_dir, "eval", iterstring) if is_main: os.makedirs(eval_dir, exist_ok=True) teacher_ckp_path = os.path.join(eval_dir, "teacher_checkpoint.pth") if isinstance(model.teacher, FSDP): state_dict_module = model.teacher elif isinstance(model, FSDP): state_dict_module = model else: state_dict_module = None state_dict_ctx = ( FSDP.state_dict_type( state_dict_module, StateDictType.FULL_STATE_DICT, state_dict_config=_FULL_STATE_DICT_CFG, ) if state_dict_module is not None else contextlib.nullcontext() ) with torch.no_grad(), state_dict_ctx: teacher_sd = model.teacher.state_dict() if is_main: torch.save({"teacher": teacher_sd}, teacher_ckp_path) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() del teacher_sd if dist.is_available() and dist.is_initialized(): dist.barrier() if not is_main: return teacher, _ = build_model_from_cfg(cfg, only_teacher=True) teacher_state = torch.load(teacher_ckp_path, map_location="cpu")["teacher"] teacher_state = {k.replace("module.", ""): v for k, v in teacher_state.items()} teacher_state = {k.replace("backbone.", ""): v for k, v in teacher_state.items() if k.startswith("backbone.")} load_msg = teacher.load_state_dict(teacher_state, strict=False) logger.info("Loaded teacher for downstream eval with msg: %s", load_msg) teacher = teacher.cuda() teacher.eval() teacher.requires_grad_(False) device = next(teacher.parameters()).device step = iteration if isinstance(iteration, int) else int(str(iteration).split("_")[-1]) def _resolve_eval_root(name, root_value): root_str = "" if root_value is None else str(root_value) if not root_str: logger.info("Skipping %s eval because %s_root is not set", name, name.lower()) return None root_path = os.path.abspath(os.path.expanduser(root_str)) if not os.path.isdir(root_path): logger.info("Skipping %s eval; dataset path missing: %s", name, root_path) return None return root_path bach_root = _resolve_eval_root("BACH", cfg.evaluation.bach_root) breakhis_root = _resolve_eval_root("BreakHis", cfg.evaluation.breakhis_root) pcam_root = _resolve_eval_root("PCam", cfg.evaluation.pcam_root) class _ResizeAndCrop(transforms.Compose): def __init__(self, size=224, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)): ops = [ transforms.Resize(size), transforms.CenterCrop(size), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std), ] super().__init__(ops) transform = _ResizeAndCrop( size=224, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ) predict_batch_size = 64 num_workers = 4 train_batch_size = 256 def _compute_embeddings(dataset, *, to_cpu=False): loader = torch.utils.data.DataLoader( dataset, batch_size=predict_batch_size, shuffle=False, num_workers=num_workers, pin_memory=True, ) feats = [] targets = [] with torch.no_grad(): for images, labels in loader: images = images.to(device, non_blocking=True) out = teacher(images, is_training=True) cls = out["x_norm_clstoken"].float() if to_cpu: feats.append(cls.cpu()) targets.append(labels.cpu()) else: labels = labels.to(device, non_blocking=True) feats.append(cls) targets.append(labels) feats = torch.cat(feats, dim=0) targets = torch.cat(targets, dim=0) return feats, targets if bach_root is not None: _BACH_TRAIN_INDEX_RANGES = [ (0, 41), (59, 60), (90, 139), (169, 240), (258, 260), (273, 345), (368, 400), ] _BACH_VAL_INDEX_RANGES = [ (41, 59), (60, 90), (139, 169), (240, 258), (260, 273), (345, 368), ] _BACH_CLASS_TO_IDX = {"Benign": 0, "InSitu": 1, "Invasive": 2, "Normal": 3} class _BACHDataset(torch.utils.data.Dataset): def __init__(self, root, split, transform): self.root = os.path.abspath(os.path.expanduser(root)) self.split = split self.transform = transform dataset_path = os.path.join(self.root, "ICIAR2018_BACH_Challenge", "Photos") self.samples = folder.make_dataset( directory=dataset_path, class_to_idx=_BACH_CLASS_TO_IDX, extensions=(".tif",), ) if len(self.samples) == 0: raise RuntimeError(f"No BACH images found in {dataset_path}") if split == "train": index_ranges = _BACH_TRAIN_INDEX_RANGES elif split == "val": index_ranges = _BACH_VAL_INDEX_RANGES else: raise ValueError("Invalid BACH split. Use 'train' or 'val'.") indices = [] for start, end in index_ranges: indices.extend(range(start, end)) self.indices = indices def __len__(self): return len(self.indices) def __getitem__(self, idx): image_path, target = self.samples[self.indices[idx]] image = Image.open(image_path).convert("RGB") if self.transform is not None: image = self.transform(image) target_tensor = torch.tensor(target, dtype=torch.long) return image, target_tensor train_ds = _BACHDataset(root=bach_root, split="train", transform=transform) val_ds = _BACHDataset(root=bach_root, split="val", transform=transform) train_feats, train_targets = _compute_embeddings(train_ds) val_feats, val_targets = _compute_embeddings(val_ds) in_features = train_feats.shape[-1] num_classes = 4 head = torch.nn.Linear(in_features, num_classes, bias=True).to(device) criterion = torch.nn.CrossEntropyLoss().to(device) optimizer = torch.optim.AdamW(head.parameters(), lr=3e-4, weight_decay=1e-2) train_dataset = torch.utils.data.TensorDataset(train_feats, train_targets) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=train_batch_size, shuffle=True, drop_last=False, ) val_dataset = torch.utils.data.TensorDataset(val_feats, val_targets) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=train_batch_size, shuffle=False, drop_last=False, ) def _eval_head(): head.eval() all_preds = [] all_targets = [] with torch.no_grad(): for feats_batch, targets_batch in val_loader: feats_batch = feats_batch.to(device, non_blocking=True) logits = head(feats_batch) preds = logits.argmax(dim=1).cpu() all_preds.append(preds) all_targets.append(targets_batch.cpu()) preds = torch.cat(all_preds, dim=0) targets = torch.cat(all_targets, dim=0) plain_acc = float((preds == targets).float().mean().item()) conf = torch.zeros(num_classes, num_classes, dtype=torch.long) indices = targets * num_classes + preds bincount = torch.bincount(indices, minlength=num_classes * num_classes) conf = bincount.view(num_classes, num_classes) per_class = conf.diag().float() / conf.sum(dim=1).clamp_min(1) balanced_acc = float(per_class.mean().item()) head.train() return plain_acc, balanced_acc max_steps = 12500 # eva uses 12500 steps with patience-based early stopping eval_every = 250 patience = 1250 steps = 0 best_plain = -1.0 best_balanced = -1.0 best_state = None steps_since_improve = 0 head.train() with tqdm(total=max_steps) as pbar: while steps < max_steps: for feats_batch, targets_batch in train_loader: feats_batch = feats_batch.to(device, non_blocking=True) targets_batch = targets_batch.to(device, non_blocking=True) logits = head(feats_batch) loss = criterion(logits, targets_batch) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() steps += 1 pbar.update(1) if steps % eval_every == 0 or steps >= max_steps: plain_acc, balanced_acc = _eval_head() if plain_acc > best_plain: best_plain = plain_acc best_balanced = balanced_acc best_state = {k: v.cpu() for k, v in head.state_dict().items()} steps_since_improve = 0 else: steps_since_improve += eval_every if steps_since_improve >= patience: steps = max_steps break if steps >= max_steps: break if best_state is not None: head.load_state_dict(best_state) bach_acc_plain, bach_acc_balanced = best_plain, best_balanced else: bach_acc_plain, bach_acc_balanced = _eval_head() logger.info( "BACH val accuracy (linear probe): plain=%.4f balanced=%.4f", bach_acc_plain, bach_acc_balanced, ) if wandb.run is not None and distributed.is_main_process(): wandb.log( { "val/BACH_BALANCED_ACCURACY": bach_acc_balanced, "val/BACH_MULTICLASS_ACCURACY": bach_acc_plain, }, step=step, ) if breakhis_root is not None: _BREAKHIS_VAL_PATIENT_IDS = { "18842D", "19979", "15275", "15792", "16875", "3909", "5287", "16716", "2773", "5695", "16184CD", "23060CD", "21998CD", "21998EF", } _BREAKHIS_CLASSES = ["TA", "MC", "F", "DC"] _BREAKHIS_CLASS_TO_IDX = {label: index for index, label in enumerate(_BREAKHIS_CLASSES)} class _BreakHisDataset(torch.utils.data.Dataset): def __init__(self, root, split, transform): self.root = os.path.abspath(os.path.expanduser(root)) self.split = split self.transform = transform dataset_path = os.path.join(self.root, "BreaKHis_v1", "histology_slides") pattern = os.path.join(dataset_path, "**", "40X", "*.png") self.image_files = sorted(glob.glob(pattern, recursive=True)) if len(self.image_files) == 0: raise RuntimeError(f"No BreakHis images found in {dataset_path}") indices = [] for idx, image_file in enumerate(self.image_files): class_name = os.path.basename(image_file).split("-")[0].split("_")[-1] if class_name not in _BREAKHIS_CLASS_TO_IDX: continue patient_id = os.path.basename(image_file).split("-")[2] if split == "train": if patient_id not in _BREAKHIS_VAL_PATIENT_IDS: indices.append(idx) elif split == "val": if patient_id in _BREAKHIS_VAL_PATIENT_IDS: indices.append(idx) else: raise ValueError("Invalid BreakHis split. Use 'train' or 'val'.") self.indices = indices def __len__(self): return len(self.indices) def __getitem__(self, idx): image_path = self.image_files[self.indices[idx]] image = Image.open(image_path).convert("RGB") if self.transform is not None: image = self.transform(image) class_name = os.path.basename(image_path).split("-")[0].split("_")[-1] target = _BREAKHIS_CLASS_TO_IDX[class_name] target_tensor = torch.tensor(target, dtype=torch.long) return image, target_tensor train_ds = _BreakHisDataset(root=breakhis_root, split="train", transform=transform) val_ds = _BreakHisDataset(root=breakhis_root, split="val", transform=transform) train_feats, train_targets = _compute_embeddings(train_ds) val_feats, val_targets = _compute_embeddings(val_ds) in_features = train_feats.shape[-1] num_classes = 4 head = torch.nn.Linear(in_features, num_classes, bias=True).to(device) criterion = torch.nn.CrossEntropyLoss().to(device) optimizer = torch.optim.AdamW(head.parameters(), lr=3e-4, weight_decay=1e-2) train_dataset = torch.utils.data.TensorDataset(train_feats, train_targets) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=train_batch_size, shuffle=True, drop_last=False, ) val_dataset = torch.utils.data.TensorDataset(val_feats, val_targets) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=train_batch_size, shuffle=False, drop_last=False, ) def _eval_head(): head.eval() all_preds = [] all_targets = [] with torch.no_grad(): for feats_batch, targets_batch in val_loader: feats_batch = feats_batch.to(device, non_blocking=True) logits = head(feats_batch) preds = logits.argmax(dim=1).cpu() all_preds.append(preds) all_targets.append(targets_batch.cpu()) preds = torch.cat(all_preds, dim=0) targets = torch.cat(all_targets, dim=0) plain_acc = float((preds == targets).float().mean().item()) conf = torch.zeros(num_classes, num_classes, dtype=torch.long) indices = targets * num_classes + preds bincount = torch.bincount(indices, minlength=num_classes * num_classes) conf = bincount.view(num_classes, num_classes) per_class = conf.diag().float() / conf.sum(dim=1).clamp_min(1) balanced_acc = float(per_class.mean().item()) head.train() return plain_acc, balanced_acc max_steps = 12500 eval_every = 250 patience = 500 steps = 0 best_plain = -1.0 best_balanced = -1.0 best_state = None steps_since_improve = 0 head.train() while steps < max_steps: for feats_batch, targets_batch in train_loader: feats_batch = feats_batch.to(device, non_blocking=True) targets_batch = targets_batch.to(device, non_blocking=True) logits = head(feats_batch) loss = criterion(logits, targets_batch) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() steps += 1 if steps % eval_every == 0 or steps >= max_steps: plain_acc, balanced_acc = _eval_head() if plain_acc > best_plain: best_plain = plain_acc best_balanced = balanced_acc best_state = {k: v.cpu() for k, v in head.state_dict().items()} steps_since_improve = 0 else: steps_since_improve += eval_every if steps_since_improve >= patience: steps = max_steps break if steps >= max_steps: break if best_state is not None: head.load_state_dict(best_state) breakhis_acc_plain, breakhis_acc_balanced = best_plain, best_balanced else: breakhis_acc_plain, breakhis_acc_balanced = _eval_head() logger.info( "BreakHis val accuracy (linear probe): plain=%.4f balanced=%.4f", breakhis_acc_plain, breakhis_acc_balanced, ) if wandb.run is not None and distributed.is_main_process(): wandb.log( { "val/BREAKHIS_BALANCED_ACCURACY": breakhis_acc_balanced, "val/BREAKHIS_MULTICLASS_ACCURACY": breakhis_acc_plain, }, step=step, ) if pcam_root is not None: def _pcam_h5_path(root, split, key): split_suffix = "valid" if split == "val" else split filename = f"camelyonpatch_level_2_split_{split_suffix}_{key}.h5" return os.path.join(root, filename) class _PCamDataset(torch.utils.data.Dataset): def __init__(self, root, split, transform): self.root = os.path.abspath(os.path.expanduser(root)) self.split = split self.transform = transform self.x_path = _pcam_h5_path(self.root, split, "x") self.y_path = _pcam_h5_path(self.root, split, "y") if not os.path.isfile(self.x_path) or not os.path.isfile(self.y_path): raise RuntimeError(f"Missing PatchCamelyon files for split {split} in {self.root}") with h5py.File(self.y_path, "r") as file: self.length = len(file["y"]) def __len__(self): return self.length def __getitem__(self, idx): with h5py.File(self.x_path, "r") as file: image_array = file["x"][idx] with h5py.File(self.y_path, "r") as file: target = file["y"][idx].squeeze() image = Image.fromarray(image_array) if image.mode != "RGB": image = image.convert("RGB") if self.transform is not None: image = self.transform(image) target_tensor = torch.tensor(float(target), dtype=torch.float32) return image, target_tensor train_ds = _PCamDataset(root=pcam_root, split="train", transform=transform) test_ds = _PCamDataset(root=pcam_root, split="test", transform=transform) num_samples_per_class = 10 with h5py.File(_pcam_h5_path(pcam_root, "train", "y"), "r") as file: targets = file["y"][:].reshape(-1) class_indices = {} for idx, target in enumerate(targets): class_idx = int(target) if class_idx not in class_indices: class_indices[class_idx] = [] class_indices[class_idx].append(idx) for class_idx, indices in class_indices.items(): if len(indices) < num_samples_per_class: raise ValueError( f"Class {class_idx} has only {len(indices)} samples, " f"which is less than the required {num_samples_per_class} samples." ) rng = np.random.default_rng(42) sampled_indices = [] for class_idx in class_indices: sampled = rng.choice( class_indices[class_idx], size=num_samples_per_class, replace=False, ).tolist() sampled_indices.extend(sampled) rng.shuffle(sampled_indices) train_subset = torch.utils.data.Subset(train_ds, sampled_indices) train_feats, train_targets = _compute_embeddings(train_subset, to_cpu=True) test_feats, test_targets = _compute_embeddings(test_ds, to_cpu=True) in_features = train_feats.shape[-1] head = torch.nn.Linear(in_features, 1, bias=True).to(device) criterion = torch.nn.BCEWithLogitsLoss().to(device) optimizer = torch.optim.AdamW(head.parameters(), lr=3e-4, weight_decay=1e-2) train_dataset = torch.utils.data.TensorDataset(train_feats, train_targets) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=train_batch_size, shuffle=False, drop_last=False, ) test_dataset = torch.utils.data.TensorDataset(test_feats, test_targets) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=train_batch_size, shuffle=False, drop_last=False, ) def _eval_head(loader): head.eval() all_preds = [] all_targets = [] with torch.no_grad(): for feats_batch, targets_batch in loader: feats_batch = feats_batch.to(device, non_blocking=True) logits = head(feats_batch).squeeze(1) preds = (logits >= 0).to(torch.long).cpu() all_preds.append(preds) all_targets.append(targets_batch.to(torch.long).cpu()) preds = torch.cat(all_preds, dim=0) targets = torch.cat(all_targets, dim=0) plain_acc = float((preds == targets).float().mean().item()) num_classes = 2 indices = targets * num_classes + preds bincount = torch.bincount(indices, minlength=num_classes * num_classes) conf = bincount.view(num_classes, num_classes) per_class = conf.diag().float() / conf.sum(dim=1).clamp_min(1) balanced_acc = float(per_class.mean().item()) head.train() return plain_acc, balanced_acc max_steps = 12500 eval_every_epochs = 10 patience_evals = 125 eval_every = eval_every_epochs * len(train_loader) patience = patience_evals * eval_every steps = 0 best_balanced = -1.0 best_state = None steps_since_improve = 0 head.train() while steps < max_steps: for feats_batch, targets_batch in train_loader: feats_batch = feats_batch.to(device, non_blocking=True) targets_batch = targets_batch.to(device, non_blocking=True) logits = head(feats_batch).squeeze(1) loss = criterion(logits, targets_batch) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() steps += 1 if steps % eval_every == 0 or steps >= max_steps: _, balanced_acc = _eval_head(test_loader) if balanced_acc > best_balanced: best_balanced = balanced_acc best_state = {k: v.cpu() for k, v in head.state_dict().items()} steps_since_improve = 0 else: steps_since_improve += eval_every if steps_since_improve >= patience: steps = max_steps break if steps >= max_steps: break if best_state is not None: head.load_state_dict(best_state) test_plain, test_balanced = _eval_head(test_loader) logger.info( "PCam test accuracy (linear probe, 10-shot): plain=%.4f balanced=%.4f", test_plain, test_balanced, ) if wandb.run is not None and distributed.is_main_process(): wandb.log( { "val/PCAM_BINARY_ACCURACY": test_plain, "val/PCAM_BALANCED_ACCURACY": test_balanced, }, step=step, ) def do_train(cfg, model, resume=False): model.train() inputs_dtype = torch.half fp16_scaler = model.fp16_scaler # for mixed precision training trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) total_params = sum(p.numel() for p in model.parameters()) if cfg.train.skip_checkpointer: print(f"\n\nSkipping FSDP checkpointer (cfg.train.skip_checkpointer={cfg.train.skip_checkpointer})\n\n") # setup optimizer optimizer = build_optimizer(cfg, model.get_params_groups()) ( lr_schedule, wd_schedule, momentum_schedule, teacher_temp_schedule, last_layer_lr_schedule, gram_weight_schedule, ) = build_schedulers(cfg) from omegaconf import OmegaConf if distributed.is_main_process(): run_id_path = Path(cfg.train.output_dir) / "wandb_run_id.txt" if resume and run_id_path.exists(): run_id = run_id_path.read_text().strip() resume_mode = "must" else: run_id_path.parent.mkdir(parents=True, exist_ok=True) run_id = wandb.util.generate_id() run_id_path.write_text(run_id) resume_mode = "allow" run = wandb.init( project="tcga-finetuning", config=OmegaConf.to_container(cfg), id=run_id, resume=resume_mode, ) repo_root = Path(__file__).resolve().parents[2] files_to_save = [ Path(__file__).resolve(), AUGMENTATION_FILE, VISION_TRANSFORMER_FILE, SSL_META_ARCH, Path(CONFIG_FILE_PATH), ] run_script = os.environ.get("DINOV2_RUN_SCRIPT") if run_script: files_to_save.append(Path(run_script).resolve()) for src in files_to_save: base_path = repo_root if src.is_relative_to(repo_root) else src.parent run.save(str(src), base_path=str(base_path), policy="now") logger.info("Trainable parameters: %s", trainable_params) logger.info("Total parameters: %s", total_params) print(f"Trainable parameters: {trainable_params}") print(f"Total parameters: {total_params}") # checkpointer if not cfg.train.skip_checkpointer: checkpointer = FSDPCheckpointer(model, cfg.train.output_dir, optimizer=optimizer, save_to_disk=True) start_iter = checkpointer.resume_or_load(cfg.MODEL.WEIGHTS, resume=resume).get("iteration", -1) + 1 else: start_iter = 0 OFFICIAL_EPOCH_LENGTH = cfg.train.OFFICIAL_EPOCH_LENGTH max_iter = cfg.optim.epochs * OFFICIAL_EPOCH_LENGTH early_stop_iter = cfg.optim.early_stop * OFFICIAL_EPOCH_LENGTH eta_target_iter = min(max_iter, early_stop_iter) if not cfg.train.skip_checkpointer: periodic_checkpointer = PeriodicCheckpointer( checkpointer, period=2 * OFFICIAL_EPOCH_LENGTH, # 2500 (was 5×=6250): 불안정 클러스터 대비 촘촘한 resume 체크포인트 max_iter=max_iter, max_to_keep=2, ) # setup data preprocessing img_size = cfg.crops.global_crops_size patch_size = cfg.student.patch_size n_tokens = (img_size // patch_size) ** 2 mask_generator = MaskingGenerator( input_size=(img_size // patch_size, img_size // patch_size), max_num_patches=0.5 * img_size // patch_size * img_size // patch_size, ) data_transform = DataAugmentationDINO( cfg.crops.global_crops_scale, cfg.crops.local_crops_scale, cfg.crops.local_crops_number, global_crops_size=cfg.crops.global_crops_size, local_crops_size=cfg.crops.local_crops_size, ) collate_fn = partial( collate_data_and_cast, mask_ratio_tuple=cfg.ibot.mask_ratio_min_max, mask_probability=cfg.ibot.mask_sample_probability, n_tokens=n_tokens, mask_generator=mask_generator, dtype=inputs_dtype, ) # setup data loader if cfg.train.streaming_from_hf: dataset_builder = partial( _build_streaming_dataset, dataset_path=str(cfg.train.streaming_dataset_path), shuffle_buffer=2000, base_seed=42, ) def decode_and_transform(item): image = Image.open(BytesIO(item["image_bytes"])) image = ImageOps.exif_transpose(image).convert("RGB") transformed = data_transform(image) slide_meta = (item["slide_path"], item["x"], item["y"], item["level"]) return (transformed, None), slide_meta class _TransformedStreamingDataset(torch.utils.data.IterableDataset): def __init__(self, dataset_builder, transform, samples_per_epoch=None, reshuffle_every=0): self._dataset_builder = dataset_builder self._transform = transform self._samples_per_epoch = samples_per_epoch self._reshuffle_every = reshuffle_every self._initialized = False self._epoch_seen = 0 self._src_iter = None def _init_or_reshuffle(self, *, force: bool = False): if force or (not self._initialized) or ( self._reshuffle_every and (self._epoch_seen % self._reshuffle_every == 0) ): src = self._dataset_builder(epoch=self._epoch_seen if self._reshuffle_every else 0) worker_info = torch.utils.data.get_worker_info() if worker_info is not None and worker_info.num_workers > 1: src = src.shard(num_shards=worker_info.num_workers, index=worker_info.id) self._src_iter = iter(src) self._initialized = True def __iter__(self): while True: self._init_or_reshuffle() # Per-RANK quota rank_quota = self._samples_per_epoch or (1 << 62) worker_info = torch.utils.data.get_worker_info() num_workers = worker_info.num_workers if worker_info is not None else 1 worker_id = worker_info.id if worker_info is not None else 0 # Split quota across workers (nearly even split) base = rank_quota // num_workers remainder = rank_quota % num_workers local_quota = base + (1 if worker_id < remainder else 0) produced = 0 while produced < local_quota: try: sample = next(self._src_iter) except StopIteration: # Refill the iterator; only reshuffle on epoch boundaries self._init_or_reshuffle(force=True) continue yield self._transform(sample) produced += 1 self._epoch_seen += 1 # Define explicit per-epoch sample budget per rank to keep ranks in lock-step samples_per_epoch = cfg.train.batch_size_per_gpu * cfg.train.OFFICIAL_EPOCH_LENGTH dataset = _TransformedStreamingDataset( dataset_builder, decode_and_transform, samples_per_epoch=samples_per_epoch, ) def _worker_init(_): torch.set_num_threads(1) os.environ.setdefault("OMP_NUM_THREADS", "1") data_loader = torch.utils.data.DataLoader( dataset, batch_size=cfg.train.batch_size_per_gpu, num_workers=cfg.train.num_workers, drop_last=True, pin_memory=True, persistent_workers=True, collate_fn=collate_fn, prefetch_factor=4, worker_init_fn=_worker_init, ) elif str(cfg.train.sample_list_path).startswith("parquet:"): # Test B: 로컬 parquet(그들 TCGA-12K)을 우리 빠른 로더로 스트리밍. from dinov2.data.openpath_wds import make_openpath_parquet_loader data_loader = make_openpath_parquet_loader( str(cfg.train.sample_list_path), batch_size=cfg.train.batch_size_per_gpu, num_workers=cfg.train.num_workers, data_transform=data_transform, collate_fn=collate_fn, prefetch_factor=cfg.train.prefetch_factor, ) elif str(cfg.train.sample_list_path).startswith("openpath:"): # OpenPath: stream pre-patched WebDataset tar shards (see openpath_wds.py). from dinov2.data.openpath_wds import make_openpath_loader data_loader = make_openpath_loader( str(cfg.train.sample_list_path), batch_size=cfg.train.batch_size_per_gpu, num_workers=cfg.train.num_workers, data_transform=data_transform, collate_fn=collate_fn, prefetch_factor=cfg.train.prefetch_factor, ) else: from dinov2.data import SamplerType, make_data_loader, make_dataset sample_list_path = str(cfg.train.sample_list_path) if not sample_list_path: raise ValueError("cfg.train.sample_list_path must be set when streaming_from_hf is False") dataset_str = f"pathology:root=/data/TCGA/:sample_list_path={sample_list_path}" dataset = make_dataset( dataset_str=dataset_str, transform=data_transform, target_transform=lambda _: (), ) sampler_type = SamplerType.SHARDED_INFINITE data_loader = make_data_loader( dataset=dataset, batch_size=cfg.train.batch_size_per_gpu, num_workers=cfg.train.num_workers, shuffle=True, seed=0, sampler_type=sampler_type, sampler_advance=0, # TODO(qas): fix this -- start_iter * cfg.train.batch_size_per_gpu, drop_last=True, collate_fn=collate_fn, persistent_workers=cfg.train.num_workers > 0, prefetch_factor=cfg.train.prefetch_factor, ) # training loop iteration = start_iter logger.info("Starting training from iteration {}".format(start_iter)) metrics_file = os.path.join(cfg.train.output_dir, "training_metrics.json") metric_logger = MetricLogger(delimiter=" ", output_file=metrics_file) header = "Training" for data in metric_logger.log_every( data_loader, 10, header, eta_target_iter + 1, start_iter, ): if iteration >= early_stop_iter: logger.info("Early stopping at iteration {}".format(iteration)) if cfg.evaluation.eval_period_iterations >= 0: do_test(cfg, model, f"training_{iteration}") torch.cuda.synchronize() if not cfg.train.skip_checkpointer: checkpointer.save(f"model_{iteration:07d}", iteration=iteration) break #Save instantly if cfg.evaluation.eval_period_iterations >= 0 and (iteration) % cfg.evaluation.eval_period_iterations == 0: do_test(cfg, model, f"training_{iteration}") torch.cuda.synchronize() if not cfg.train.skip_checkpointer: periodic_checkpointer.step(iteration) current_batch_size = data["collated_global_crops"].shape[0] / 2 if iteration > max_iter: return nan_mask = torch.isnan(data["collated_global_crops"]) nan_mask2 = torch.isnan(data["collated_local_crops"]) if nan_mask.any(): print("found nan in input data") print(data[indexes]) # apply schedules lr = lr_schedule[iteration] wd = wd_schedule[iteration] mom = momentum_schedule[iteration] teacher_temp = teacher_temp_schedule[iteration] last_layer_lr = last_layer_lr_schedule[iteration] gram_weight = gram_weight_schedule[iteration] apply_optim_scheduler(optimizer, lr, wd, last_layer_lr) # compute losses optimizer.zero_grad(set_to_none=True) loss_dict = model.forward_backward(data, teacher_temp=teacher_temp, gram_weight=gram_weight) # clip gradients if fp16_scaler is not None: if cfg.optim.clip_grad: fp16_scaler.unscale_(optimizer) for v in model.student.values(): v.clip_grad_norm_(cfg.optim.clip_grad) fp16_scaler.step(optimizer) fp16_scaler.update() else: if cfg.optim.clip_grad: for v in model.student.values(): v.clip_grad_norm_(cfg.optim.clip_grad) optimizer.step() # perform teacher EMA update model.update_teacher(mom) # logging if distributed.get_global_size() > 1: for v in loss_dict.values(): torch.distributed.all_reduce(v) loss_dict_reduced = {k: v.item() / distributed.get_global_size() for k, v in loss_dict.items()} if math.isnan(sum(loss_dict_reduced.values())): print(sum(loss_dict_reduced.values())) logger.info("NaN detected") print(data["indexes"]) for name, param in model.named_parameters(): if torch.isnan(param.data).any(): print(f"NaNs found in parameter: {name}") raise AssertionError losses_reduced = sum(loss for loss in loss_dict_reduced.values()) metric_logger.update(lr=lr) metric_logger.update(wd=wd) metric_logger.update(mom=mom) metric_logger.update(last_layer_lr=last_layer_lr) metric_logger.update(current_batch_size=current_batch_size) metric_logger.update(total_loss=losses_reduced, **loss_dict_reduced) if distributed.is_main_process(): scalar_logs = { "Learning Rate": lr, "Momentum": mom, "Last Layer LR": last_layer_lr, "Total Loss": losses_reduced, } wandb.log({**scalar_logs, **loss_dict_reduced}, step=iteration) # Synchronize the GPU to ensure all operations are complete before measuring torch.cuda.synchronize() iteration = iteration + 1 metric_logger.synchronize_between_processes() return {k: meter.global_avg for k, meter in metric_logger.meters.items()} def main(args): cfg = setup(args) print(cfg) model = SSLMetaArch(cfg).to(torch.device("cuda")) #Load model here from pretrained. if cfg.train.use_pretrained: _load_pretrained_backbone(cfg, model) _freeze_student_backbone_except_last_n(cfg, model) model.prepare_for_distributed_training() logger.info("Model:\n{}".format(model)) if args.eval_only and not cfg.train.skip_checkpointer: iteration = ( FSDPCheckpointer(model, save_dir=cfg.train.output_dir) .resume_or_load(cfg.MODEL.WEIGHTS, resume=not args.no_resume) .get("iteration", -1) + 1 ) return do_test(cfg, model, f"manual_{iteration}") do_train(cfg, model, resume=not args.no_resume) if __name__ == "__main__": args = get_args_parser(add_help=True).parse_args() if not args.config_file: raise ValueError("config file path must be provided") CONFIG_FILE_PATH = os.path.abspath(args.config_file) main(args)