YannAhlgrim
724 lines
26 KiB
Python
724 lines
26 KiB
Python
import os
|
|
import random
|
|
import resource
|
|
import shutil
|
|
import sys
|
|
import json
|
|
import time
|
|
import yaml
|
|
import logging
|
|
|
|
import numpy as np
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.distributed as dist
|
|
from torch.nn.parallel import DistributedDataParallel
|
|
|
|
from src.datasets.wilds import make_iwildcam
|
|
from src.helper import init_model
|
|
from src.models.head import ViTClassifier
|
|
from src.transforms import make_transforms, make_transform_eval
|
|
from src.utils.distributed import init_distributed
|
|
from src.utils.logging import CSVLogger, AverageMeter, resolve_log_dir
|
|
from src.utils.optimizers import LARS
|
|
from src.eval_wilds import main as eval_wilds_main
|
|
|
|
# --
|
|
log_freq = 10
|
|
checkpoint_freq = 50
|
|
# --
|
|
|
|
_GLOBAL_SEED = 0
|
|
np.random.seed(_GLOBAL_SEED)
|
|
torch.manual_seed(_GLOBAL_SEED)
|
|
torch.backends.cudnn.benchmark = True
|
|
|
|
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
|
|
logger = logging.getLogger()
|
|
|
|
|
|
def _set_seed(seed):
|
|
"""Seed all RNGs so multi-seed runs are reproducible and distinct."""
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
|
|
def _format_hms(seconds):
|
|
seconds = int(round(seconds))
|
|
h, rem = divmod(seconds, 3600)
|
|
m, s = divmod(rem, 60)
|
|
return f"{h:02d}:{m:02d}:{s:02d}"
|
|
|
|
|
|
def _peak_host_ram_gb():
|
|
"""Peak resident set size (RSS) of this process, in GB.
|
|
|
|
Uses resource.getrusage(RUSAGE_SELF).ru_maxrss, which on Linux is reported
|
|
in kilobytes. This is the process high-water mark; with tasks_per_node=1 it
|
|
reflects the whole training worker. Compare against the SLURM mem request
|
|
(e.g. 180G) to right-size future jobs.
|
|
"""
|
|
try:
|
|
maxrss_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
|
|
return float(maxrss_kb) / (1024.0 * 1024.0)
|
|
except (ValueError, OSError):
|
|
return None
|
|
|
|
|
|
def _peak_gpu_mem_gb(device):
|
|
"""Peak allocated and reserved GPU memory (GB) since the last reset."""
|
|
if not torch.cuda.is_available():
|
|
return None, None
|
|
try:
|
|
alloc = torch.cuda.max_memory_allocated(device) / 1e9
|
|
reserved = torch.cuda.max_memory_reserved(device) / 1e9
|
|
return float(alloc), float(reserved)
|
|
except (RuntimeError, ValueError):
|
|
return None, None
|
|
|
|
|
|
def strip_module_prefix(state_dict):
|
|
if not any(k.startswith("module.") for k in state_dict.keys()):
|
|
return state_dict
|
|
return {
|
|
k[len("module.") :] if k.startswith("module.") else k: v
|
|
for k, v in state_dict.items()
|
|
}
|
|
|
|
|
|
def distributed_average(value, device):
|
|
tensor = torch.tensor([value], device=device, dtype=torch.float64)
|
|
if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
|
|
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
|
|
tensor /= dist.get_world_size()
|
|
return float(tensor.item())
|
|
|
|
|
|
def distributed_sum(value, device):
|
|
tensor = torch.tensor([value], device=device, dtype=torch.float64)
|
|
if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
|
|
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
|
|
return float(tensor.item())
|
|
|
|
|
|
def _find_metric(metrics, key):
|
|
if isinstance(metrics, dict):
|
|
if key in metrics:
|
|
return metrics[key]
|
|
for value in metrics.values():
|
|
found = _find_metric(value, key)
|
|
if found is not None:
|
|
return found
|
|
if isinstance(metrics, list):
|
|
for item in metrics:
|
|
found = _find_metric(item, key)
|
|
if found is not None:
|
|
return found
|
|
return None
|
|
|
|
|
|
def _collect_eval_rows(eval_root, metric_key):
|
|
rows = []
|
|
for dirpath, _, filenames in os.walk(eval_root):
|
|
for fname in filenames:
|
|
if not fname.endswith("_metrics.json"):
|
|
continue
|
|
path = os.path.join(dirpath, fname)
|
|
try:
|
|
with open(path, "r") as f:
|
|
metrics = json.load(f)
|
|
except (OSError, json.JSONDecodeError):
|
|
continue
|
|
value = _find_metric(metrics, metric_key)
|
|
if value is None:
|
|
continue
|
|
run_name = os.path.basename(os.path.dirname(path))
|
|
rows.append((float(value), run_name, path))
|
|
return rows
|
|
|
|
|
|
class EarlyStopping:
|
|
def __init__(
|
|
self,
|
|
enabled=True,
|
|
patience=10,
|
|
min_delta=0.0,
|
|
min_epochs=0,
|
|
restore_best_weights=True,
|
|
):
|
|
self.enabled = enabled
|
|
self.patience = patience
|
|
self.min_delta = min_delta
|
|
self.min_epochs = min_epochs
|
|
self.restore_best_weights = restore_best_weights
|
|
self.best_metric = float("inf")
|
|
self.best_epoch = -1
|
|
self.bad_epochs = 0
|
|
self.best_state = None
|
|
|
|
def state_dict(self):
|
|
return {
|
|
"enabled": self.enabled,
|
|
"patience": self.patience,
|
|
"min_delta": self.min_delta,
|
|
"min_epochs": self.min_epochs,
|
|
"restore_best_weights": self.restore_best_weights,
|
|
"best_metric": self.best_metric,
|
|
"best_epoch": self.best_epoch,
|
|
"bad_epochs": self.bad_epochs,
|
|
}
|
|
|
|
def load_state_dict(self, state):
|
|
if not state:
|
|
return
|
|
self.enabled = state.get("enabled", self.enabled)
|
|
self.patience = state.get("patience", self.patience)
|
|
self.min_delta = state.get("min_delta", self.min_delta)
|
|
self.min_epochs = state.get("min_epochs", self.min_epochs)
|
|
self.restore_best_weights = state.get(
|
|
"restore_best_weights", self.restore_best_weights
|
|
)
|
|
self.best_metric = state.get("best_metric", self.best_metric)
|
|
self.best_epoch = state.get("best_epoch", self.best_epoch)
|
|
self.bad_epochs = state.get("bad_epochs", self.bad_epochs)
|
|
|
|
def step(self, epoch, metric, model_module):
|
|
if not self.enabled:
|
|
return False, False
|
|
|
|
improved = metric < (self.best_metric - self.min_delta)
|
|
if improved:
|
|
self.best_metric = metric
|
|
self.best_epoch = epoch
|
|
self.bad_epochs = 0
|
|
if self.restore_best_weights:
|
|
self.best_state = {
|
|
k: v.detach().cpu().clone()
|
|
for k, v in model_module.state_dict().items()
|
|
}
|
|
return True, False
|
|
|
|
self.bad_epochs += 1
|
|
should_stop = (
|
|
epoch + 1
|
|
) >= self.min_epochs and self.bad_epochs >= self.patience
|
|
return False, should_stop
|
|
|
|
def restore(self, model_module, device):
|
|
if self.restore_best_weights and self.best_state is not None:
|
|
model_module.load_state_dict(self.best_state)
|
|
model_module.to(device)
|
|
|
|
|
|
def evaluate(model, loader, criterion, device, use_bfloat16):
|
|
model.eval()
|
|
loss_sum = 0.0
|
|
n_correct = 0.0
|
|
n_total = 0.0
|
|
|
|
with torch.no_grad():
|
|
for imgs, labels in loader:
|
|
imgs = imgs.to(device, non_blocking=True)
|
|
labels = labels.to(device, non_blocking=True)
|
|
|
|
with torch.cuda.amp.autocast(enabled=use_bfloat16, dtype=torch.bfloat16):
|
|
outputs = model(imgs)
|
|
loss = criterion(outputs, labels)
|
|
|
|
batch_size = labels.size(0)
|
|
preds = outputs.argmax(dim=1)
|
|
n_correct += float((preds == labels).sum().item())
|
|
n_total += float(batch_size)
|
|
loss_sum += float(loss.item()) * float(batch_size)
|
|
|
|
global_loss_sum = distributed_sum(loss_sum, device)
|
|
global_correct = distributed_sum(n_correct, device)
|
|
global_total = distributed_sum(n_total, device)
|
|
|
|
val_loss = global_loss_sum / max(global_total, 1.0)
|
|
val_acc = global_correct / max(global_total, 1.0)
|
|
return val_loss, val_acc
|
|
|
|
|
|
def main(args, resume_preempt=False):
|
|
del resume_preempt
|
|
|
|
world_size, rank = init_distributed()
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA is required for supervised distributed training")
|
|
device = torch.device(f"cuda:{torch.cuda.current_device()}")
|
|
|
|
m_args = args["meta"]
|
|
o_args = args["optimization"]
|
|
d_args = args["data"]
|
|
mk_args = args["mask"]
|
|
l_args = args["logging"]
|
|
v_args = args["validation"]
|
|
es_args = o_args["early_stopping"]
|
|
|
|
seed = int(m_args.get("seed", _GLOBAL_SEED))
|
|
_set_seed(seed)
|
|
logger.info(f"Using seed {seed}")
|
|
accum_steps = o_args.get("gradient_accumulation_steps", 1)
|
|
use_gradient_checkpointing = m_args.get("use_gradient_checkpointing", True)
|
|
|
|
folder = resolve_log_dir(args, stage="train")
|
|
tag = l_args["write_tag"]
|
|
|
|
with open(os.path.join(folder, "params-supervised.yaml"), "w") as f:
|
|
yaml.dump(args, f)
|
|
with open(os.path.join(folder, "params.yaml"), "w") as f:
|
|
yaml.dump(args, f)
|
|
|
|
save_path = os.path.join(folder, f"{tag}" + "-ep{epoch}.pth.tar")
|
|
latest_path = os.path.join(folder, f"{tag}-latest.pth.tar")
|
|
best_path = os.path.join(folder, f"{tag}-best.pth.tar")
|
|
log_file = os.path.join(folder, f"{tag}_r{rank}.csv")
|
|
|
|
csv_logger = CSVLogger(
|
|
log_file,
|
|
("%d", "epoch"),
|
|
("%.6f", "train_loss"),
|
|
("%.6f", "val_loss"),
|
|
("%.6f", "val_acc"),
|
|
("%.6e", "lr"),
|
|
("%.6f", "best_val_loss"),
|
|
("%d", "best_epoch"),
|
|
("%d", "early_stop"),
|
|
)
|
|
|
|
encoder, _ = init_model(
|
|
device=device,
|
|
patch_size=mk_args["patch_size"],
|
|
crop_size=d_args["crop_size"],
|
|
model_name=m_args["model_name"],
|
|
use_gradient_checkpointing=use_gradient_checkpointing,
|
|
)
|
|
|
|
representation_type = m_args.get("representation_type", "last_avgpool")
|
|
head_type = m_args.get("head_type", "linear")
|
|
model = ViTClassifier(
|
|
encoder,
|
|
m_args["num_classes"],
|
|
m_args["embed_dim"],
|
|
representation_type=representation_type,
|
|
head_type=head_type,
|
|
).to(device)
|
|
|
|
if o_args["freeze_weights"]:
|
|
logger.info("Freezing encoder weights (Linear Probing mode)")
|
|
for param in model.encoder.parameters():
|
|
param.requires_grad = False
|
|
model.encoder.eval()
|
|
else:
|
|
logger.info("Training full model (Fine-tuning mode)")
|
|
|
|
params = [p for p in model.parameters() if p.requires_grad]
|
|
optimizer_name = o_args["optimizer"].lower()
|
|
if optimizer_name == "adamw":
|
|
optimizer = torch.optim.AdamW(
|
|
params, lr=o_args["lr"], weight_decay=o_args["weight_decay"]
|
|
)
|
|
elif optimizer_name == "lars":
|
|
optimizer = LARS(
|
|
params,
|
|
lr=o_args["lr"],
|
|
weight_decay=o_args["weight_decay"],
|
|
momentum=o_args.get("momentum", 0.9),
|
|
eta=o_args.get("lars_eta", 0.001),
|
|
eps=o_args.get("lars_eps", 1e-8),
|
|
exclude_bias_and_norm=o_args.get("lars_exclude_bias_and_norm", True),
|
|
)
|
|
else:
|
|
optimizer = torch.optim.SGD(
|
|
params,
|
|
lr=o_args["lr"],
|
|
momentum=o_args.get("momentum", 0.9),
|
|
weight_decay=o_args["weight_decay"],
|
|
)
|
|
|
|
scheduler = None
|
|
lr_schedule = o_args.get("lr_schedule", "cosine").lower()
|
|
if lr_schedule == "step":
|
|
scheduler = torch.optim.lr_scheduler.MultiStepLR(
|
|
optimizer,
|
|
milestones=o_args.get("step_milestones", [15, 30, 45]),
|
|
gamma=o_args.get("step_gamma", 0.1),
|
|
)
|
|
elif lr_schedule == "cosine":
|
|
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
|
|
optimizer, T_max=o_args["epochs"], eta_min=o_args["final_lr"]
|
|
)
|
|
|
|
train_transform = make_transforms(
|
|
crop_size=d_args["crop_size"],
|
|
crop_scale=tuple(d_args["crop_scale"]),
|
|
horizontal_flip=d_args["use_horizontal_flip"],
|
|
color_distortion=d_args["use_color_distortion"],
|
|
color_jitter=d_args["color_jitter_strength"],
|
|
gaussian_blur=d_args["use_gaussian_blur"],
|
|
use_random_resized_crop=d_args.get("use_random_resized_crop", True),
|
|
)
|
|
val_transform = make_transform_eval(
|
|
crop_size=d_args["crop_size"],
|
|
)
|
|
|
|
_, train_loader, train_sampler = make_iwildcam(
|
|
transform=train_transform,
|
|
split="train",
|
|
batch_size=d_args["batch_size"],
|
|
root_path=d_args["root_path"],
|
|
rank=rank,
|
|
world_size=world_size,
|
|
collator=None,
|
|
num_workers=d_args["num_workers"],
|
|
pin_mem=d_args["pin_mem"],
|
|
drop_last=True,
|
|
)
|
|
|
|
_, val_loader, val_sampler = make_iwildcam(
|
|
transform=val_transform,
|
|
split="val",
|
|
batch_size=d_args["batch_size"],
|
|
root_path=d_args["root_path"],
|
|
rank=rank,
|
|
world_size=world_size,
|
|
collator=None,
|
|
num_workers=d_args["num_workers"],
|
|
pin_mem=d_args["pin_mem"],
|
|
drop_last=False,
|
|
)
|
|
|
|
criterion = nn.CrossEntropyLoss().to(device)
|
|
model = DistributedDataParallel(model, device_ids=[torch.cuda.current_device()])
|
|
|
|
eval_every = int(v_args["eval_every"])
|
|
|
|
early_stopper = EarlyStopping(
|
|
enabled=es_args["enabled"],
|
|
patience=es_args["patience"],
|
|
min_delta=es_args["min_delta"],
|
|
min_epochs=es_args["min_epochs"],
|
|
restore_best_weights=es_args["restore_best_weights"],
|
|
)
|
|
|
|
start_epoch = 0
|
|
|
|
checkpoint_to_load = None
|
|
resuming_interrupted = False
|
|
if os.path.exists(latest_path):
|
|
checkpoint_to_load = latest_path
|
|
resuming_interrupted = True
|
|
elif m_args["load_checkpoint"]:
|
|
r_file = m_args["read_checkpoint"]
|
|
checkpoint_folder = m_args["checkpoint_folder"]
|
|
if os.path.isabs(r_file):
|
|
checkpoint_to_load = r_file
|
|
else:
|
|
checkpoint_to_load = os.path.join(checkpoint_folder, r_file)
|
|
|
|
if checkpoint_to_load is not None and not os.path.exists(checkpoint_to_load):
|
|
raise FileNotFoundError(f"Checkpoint not found: {checkpoint_to_load}")
|
|
|
|
if checkpoint_to_load and os.path.exists(checkpoint_to_load):
|
|
checkpoint = torch.load(checkpoint_to_load, map_location="cpu")
|
|
if resuming_interrupted and "model" in checkpoint:
|
|
model.module.load_state_dict(checkpoint["model"])
|
|
if "opt" in checkpoint:
|
|
optimizer.load_state_dict(checkpoint["opt"])
|
|
if scheduler is not None and "scheduler" in checkpoint:
|
|
scheduler.load_state_dict(checkpoint["scheduler"])
|
|
start_epoch = int(checkpoint.get("epoch", 0))
|
|
early_stopper.load_state_dict(checkpoint.get("early_stopping", {}))
|
|
logger.info(
|
|
f"Resuming training from {checkpoint_to_load} at epoch {start_epoch}"
|
|
)
|
|
else:
|
|
encoder_state = checkpoint.get("encoder")
|
|
if encoder_state is None and "model" in checkpoint:
|
|
encoder_state = {
|
|
k.replace("encoder.", "", 1): v
|
|
for k, v in checkpoint["model"].items()
|
|
if k.startswith("encoder.")
|
|
}
|
|
if encoder_state is None:
|
|
raise KeyError(
|
|
f"No encoder weights found in checkpoint: {checkpoint_to_load}"
|
|
)
|
|
encoder_state = strip_module_prefix(encoder_state)
|
|
msg = model.module.encoder.load_state_dict(encoder_state, strict=False)
|
|
logger.info(
|
|
f"Loaded pre-trained encoder from {checkpoint_to_load} with msg: {msg}"
|
|
)
|
|
|
|
def save_checkpoint(epoch, train_loss, val_loss, val_acc, is_best=False):
|
|
save_dict = {
|
|
"model": model.module.state_dict(),
|
|
"opt": optimizer.state_dict(),
|
|
"scheduler": None if scheduler is None else scheduler.state_dict(),
|
|
"epoch": epoch,
|
|
"train_loss": train_loss,
|
|
"val_loss": val_loss,
|
|
"val_acc": val_acc,
|
|
"args": args,
|
|
"early_stopping": early_stopper.state_dict(),
|
|
}
|
|
if rank == 0:
|
|
torch.save(save_dict, latest_path)
|
|
if epoch % checkpoint_freq == 0:
|
|
torch.save(save_dict, save_path.format(epoch=epoch))
|
|
if is_best:
|
|
torch.save(save_dict, best_path)
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
|
|
train_start_time = time.perf_counter()
|
|
epochs_run = 0
|
|
early_stopped = False
|
|
for epoch in range(start_epoch, o_args["epochs"]):
|
|
train_sampler.set_epoch(epoch)
|
|
val_sampler.set_epoch(epoch)
|
|
epochs_run = epoch + 1
|
|
|
|
model.train()
|
|
if o_args["freeze_weights"]:
|
|
model.module.encoder.eval()
|
|
|
|
loss_meter = AverageMeter()
|
|
|
|
optimizer.zero_grad()
|
|
for itr, (imgs, labels) in enumerate(train_loader):
|
|
imgs = imgs.to(device, non_blocking=True)
|
|
labels = labels.to(device, non_blocking=True)
|
|
|
|
with torch.cuda.amp.autocast(
|
|
enabled=m_args["use_bfloat16"], dtype=torch.bfloat16
|
|
):
|
|
outputs = model(imgs)
|
|
loss = criterion(outputs, labels)
|
|
|
|
(loss / accum_steps).backward()
|
|
|
|
if (itr + 1) % accum_steps == 0:
|
|
optimizer.step()
|
|
optimizer.zero_grad()
|
|
|
|
loss_meter.update(loss.item(), n=labels.size(0))
|
|
|
|
if itr % log_freq == 0 and rank == 0:
|
|
logger.info(
|
|
f"Epoch {epoch + 1} [{itr}/{len(train_loader)}] Train Loss: {loss_meter.avg:.4f}"
|
|
)
|
|
|
|
train_loss = distributed_average(loss_meter.avg, device)
|
|
|
|
do_eval = ((epoch + 1) % eval_every == 0) or (epoch + 1 == o_args["epochs"])
|
|
if do_eval:
|
|
val_loss, val_acc = evaluate(
|
|
model=model,
|
|
loader=val_loader,
|
|
criterion=criterion,
|
|
device=device,
|
|
use_bfloat16=m_args["use_bfloat16"],
|
|
)
|
|
is_best, should_stop = early_stopper.step(epoch + 1, val_loss, model.module)
|
|
else:
|
|
val_loss = float("nan")
|
|
val_acc = float("nan")
|
|
is_best, should_stop = False, False
|
|
|
|
if scheduler is not None:
|
|
scheduler.step()
|
|
|
|
if rank == 0:
|
|
logger.info(
|
|
f"Epoch {epoch + 1} done | train_loss={train_loss:.6f} val_loss={val_loss:.6f} val_acc={val_acc:.6f} best_val_loss={early_stopper.best_metric:.6f}"
|
|
)
|
|
|
|
csv_logger.log(
|
|
epoch + 1,
|
|
train_loss,
|
|
val_loss,
|
|
val_acc,
|
|
optimizer.param_groups[0]["lr"],
|
|
early_stopper.best_metric,
|
|
early_stopper.best_epoch,
|
|
int(should_stop),
|
|
)
|
|
|
|
save_checkpoint(epoch + 1, train_loss, val_loss, val_acc, is_best=is_best)
|
|
|
|
stop_tensor = torch.tensor([int(should_stop)], device=device)
|
|
if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
|
|
dist.broadcast(stop_tensor, src=0)
|
|
|
|
if bool(stop_tensor.item()):
|
|
early_stopped = True
|
|
if rank == 0:
|
|
logger.info(
|
|
f"Early stopping at epoch {epoch + 1}. Best val_loss={early_stopper.best_metric:.6f} @ epoch {early_stopper.best_epoch}"
|
|
)
|
|
break
|
|
|
|
train_time_seconds = time.perf_counter() - train_start_time
|
|
peak_host_ram_gb = _peak_host_ram_gb()
|
|
peak_gpu_alloc_gb, peak_gpu_reserved_gb = _peak_gpu_mem_gb(device)
|
|
if rank == 0:
|
|
logger.info(
|
|
f"Training finished: epochs_run={epochs_run} "
|
|
f"early_stopped={early_stopped} "
|
|
f"train_time={_format_hms(train_time_seconds)} "
|
|
f"({train_time_seconds:.1f}s) "
|
|
f"peak_host_ram_gb={peak_host_ram_gb} "
|
|
f"peak_gpu_alloc_gb={peak_gpu_alloc_gb} "
|
|
f"peak_gpu_reserved_gb={peak_gpu_reserved_gb}"
|
|
)
|
|
|
|
if early_stopper.enabled and early_stopper.restore_best_weights:
|
|
if rank == 0:
|
|
logger.info("Restoring best model weights before exit")
|
|
early_stopper.restore(model.module, device)
|
|
if rank == 0:
|
|
torch.save(
|
|
{
|
|
"model": model.module.state_dict(),
|
|
"epoch": early_stopper.best_epoch,
|
|
"val_loss": early_stopper.best_metric,
|
|
"args": args,
|
|
},
|
|
best_path,
|
|
)
|
|
|
|
if rank == 0:
|
|
# Evaluate on both the in-distribution (id_test) and out-of-distribution
|
|
# (test) splits so the generalization gap can be measured.
|
|
# WILDS-iWildCam: "id_test" == Test ID, "test" == Test OOD.
|
|
eval_splits = [
|
|
("id_test", "iwildcam_id_test"),
|
|
("test", "iwildcam_test"),
|
|
]
|
|
|
|
def _make_eval_args(split, write_tag):
|
|
return {
|
|
"meta": {
|
|
"seed": seed,
|
|
"model_name": m_args["model_name"],
|
|
"embed_dim": m_args["embed_dim"],
|
|
"num_classes": m_args["num_classes"],
|
|
"patch_size": mk_args.get("patch_size", m_args.get("patch_size", 16)),
|
|
"crop_size": d_args.get("crop_size", m_args.get("crop_size", 224)),
|
|
"use_bfloat16": m_args.get("use_bfloat16", True),
|
|
"representation_type": representation_type,
|
|
"head_type": head_type,
|
|
"checkpoint_path": best_path,
|
|
"force_single_process": True,
|
|
},
|
|
"data": {
|
|
"batch_size": d_args.get("batch_size", 128),
|
|
"root_path": d_args.get("root_path", "./wilds_data"),
|
|
"num_workers": d_args.get("num_workers", 8),
|
|
"pin_mem": d_args.get("pin_mem", True),
|
|
"split": split,
|
|
"download": True,
|
|
},
|
|
"logging": {
|
|
"write_tag": write_tag,
|
|
"auto_folder": True,
|
|
},
|
|
}
|
|
|
|
eval_results = {}
|
|
eval_folder = None
|
|
for split, write_tag in eval_splits:
|
|
eval_args = _make_eval_args(split, write_tag)
|
|
result = eval_wilds_main(args=eval_args)
|
|
eval_results[split] = result
|
|
if eval_folder is None:
|
|
eval_folder = eval_args.get("logging", {}).get("folder")
|
|
|
|
# Common run-level info folded into every metrics JSON + the params summary.
|
|
run_info = {
|
|
"seed": seed,
|
|
"train_time_seconds": float(train_time_seconds),
|
|
"train_time_hms": _format_hms(train_time_seconds),
|
|
"epochs_run": int(epochs_run),
|
|
"configured_epochs": int(o_args["epochs"]),
|
|
"best_epoch": int(early_stopper.best_epoch),
|
|
"early_stopped": bool(early_stopped),
|
|
"best_val_loss": float(early_stopper.best_metric),
|
|
"peak_host_ram_gb": peak_host_ram_gb,
|
|
"peak_gpu_alloc_gb": peak_gpu_alloc_gb,
|
|
"peak_gpu_reserved_gb": peak_gpu_reserved_gb,
|
|
}
|
|
|
|
# Fold run_info into each split's metrics JSON so an aggregator can read
|
|
# metrics + seed + timing + epochs from a single file per split.
|
|
for split, result in eval_results.items():
|
|
if not result:
|
|
continue
|
|
metrics_path = result.get("metrics_path")
|
|
if not metrics_path or not os.path.exists(metrics_path):
|
|
continue
|
|
try:
|
|
with open(metrics_path, "r") as f:
|
|
metrics_obj = json.load(f)
|
|
metrics_obj["run_info"] = run_info
|
|
metrics_obj["split"] = split
|
|
with open(metrics_path, "w") as f:
|
|
json.dump(metrics_obj, f, indent=2, sort_keys=True)
|
|
except (OSError, json.JSONDecodeError):
|
|
logger.warning(f"Could not augment metrics JSON for split {split}")
|
|
|
|
if eval_folder:
|
|
try:
|
|
params_out = yaml.safe_load(yaml.dump(args))
|
|
params_out.setdefault("meta", {})["representation_type"] = representation_type
|
|
params_out.setdefault("meta", {})["head_type"] = head_type
|
|
params_out["results"] = {
|
|
"best_val_loss": float(early_stopper.best_metric),
|
|
"best_epoch": int(early_stopper.best_epoch),
|
|
"best_checkpoint": best_path,
|
|
"seed": seed,
|
|
"train_time_seconds": float(train_time_seconds),
|
|
"train_time_hms": _format_hms(train_time_seconds),
|
|
"epochs_run": int(epochs_run),
|
|
"configured_epochs": int(o_args["epochs"]),
|
|
"early_stopped": bool(early_stopped),
|
|
"peak_host_ram_gb": peak_host_ram_gb,
|
|
"peak_gpu_alloc_gb": peak_gpu_alloc_gb,
|
|
"peak_gpu_reserved_gb": peak_gpu_reserved_gb,
|
|
"eval_metrics_id_test": (
|
|
eval_results.get("id_test", {}).get("metrics")
|
|
if eval_results.get("id_test")
|
|
else None
|
|
),
|
|
"eval_metrics_test": (
|
|
eval_results.get("test", {}).get("metrics")
|
|
if eval_results.get("test")
|
|
else None
|
|
),
|
|
}
|
|
with open(os.path.join(eval_folder, "params-supervised.yaml"), "w") as f:
|
|
yaml.dump(params_out, f)
|
|
with open(os.path.join(eval_folder, "params.yaml"), "w") as f:
|
|
yaml.dump(params_out, f)
|
|
except OSError:
|
|
logger.warning("Could not write supervised params to eval folder")
|
|
|
|
if os.path.exists(folder):
|
|
try:
|
|
shutil.rmtree(folder)
|
|
except OSError:
|
|
logger.warning("Could not remove supervised run folder")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise RuntimeError(
|
|
"Use main_distributed_supervised.py to launch this script with a config file."
|
|
)
|