YannAhlgrim
213 lines
6.4 KiB
Python
213 lines
6.4 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import yaml
|
|
import logging
|
|
|
|
import numpy as np
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from src.datasets.wilds import make_iwildcam_eval
|
|
from src.helper import init_model
|
|
from src.models.head import ViTClassifier
|
|
from src.transforms import make_transform_eval
|
|
from src.utils.distributed import init_distributed
|
|
from src.utils.logging import resolve_log_dir
|
|
|
|
|
|
_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 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 _load_yaml(path):
|
|
with open(path, "r") as f:
|
|
return yaml.load(f, Loader=yaml.FullLoader)
|
|
|
|
|
|
def _get_seed(args):
|
|
return int(args.get("meta", {}).get("seed", _GLOBAL_SEED))
|
|
|
|
|
|
def _set_seed(seed):
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
|
|
|
|
def _resolve_checkpoint_path(meta_args):
|
|
if meta_args.get("checkpoint_path"):
|
|
return meta_args["checkpoint_path"]
|
|
folder = meta_args.get("checkpoint_folder")
|
|
fname = meta_args.get("read_checkpoint")
|
|
if folder is None or fname is None:
|
|
return None
|
|
return fname if os.path.isabs(fname) else os.path.join(folder, fname)
|
|
|
|
|
|
def _load_model_state(model, checkpoint_path, device):
|
|
if checkpoint_path is None:
|
|
raise ValueError("checkpoint_path or checkpoint_folder/read_checkpoint required")
|
|
if not os.path.exists(checkpoint_path):
|
|
raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
|
|
|
|
checkpoint = torch.load(checkpoint_path, map_location="cpu")
|
|
if "model" not in checkpoint:
|
|
raise KeyError(f"No model weights found in checkpoint: {checkpoint_path}")
|
|
state = strip_module_prefix(checkpoint["model"])
|
|
msg = model.load_state_dict(state, strict=True)
|
|
model.to(device)
|
|
logger.info(f"Loaded model from {checkpoint_path} with msg: {msg}")
|
|
|
|
|
|
def main(args):
|
|
force_single = bool(args.get("meta", {}).get("force_single_process", False))
|
|
if force_single:
|
|
world_size, rank = 1, 0
|
|
else:
|
|
world_size, rank = init_distributed()
|
|
|
|
if (
|
|
not force_single
|
|
and dist.is_available()
|
|
and dist.is_initialized()
|
|
and world_size > 1
|
|
and rank != 0
|
|
):
|
|
dist.barrier()
|
|
dist.destroy_process_group()
|
|
return
|
|
|
|
seed = _get_seed(args)
|
|
_set_seed(seed)
|
|
|
|
meta_args = args.get("meta", {})
|
|
data_args = args.get("data", {})
|
|
log_args = args.get("logging", {})
|
|
|
|
if not torch.cuda.is_available():
|
|
device = torch.device("cpu")
|
|
else:
|
|
device = torch.device(f"cuda:{torch.cuda.current_device()}")
|
|
|
|
folder = resolve_log_dir(args, stage="eval")
|
|
tag = log_args.get("write_tag", "wilds_eval")
|
|
|
|
params_path = os.path.join(folder, "params-eval.yaml")
|
|
with open(params_path, "w") as f:
|
|
yaml.dump(args, f)
|
|
with open(os.path.join(folder, "params.yaml"), "w") as f:
|
|
yaml.dump(args, f)
|
|
|
|
model_name = meta_args["model_name"]
|
|
embed_dim = meta_args["embed_dim"]
|
|
num_classes = meta_args["num_classes"]
|
|
patch_size = meta_args.get("patch_size", 16)
|
|
crop_size = meta_args.get("crop_size", 224)
|
|
use_bfloat16 = bool(meta_args.get("use_bfloat16", True))
|
|
use_autocast = use_bfloat16 and torch.cuda.is_available()
|
|
|
|
encoder, _ = init_model(
|
|
device=device,
|
|
patch_size=patch_size,
|
|
crop_size=crop_size,
|
|
model_name=model_name,
|
|
)
|
|
model = ViTClassifier(
|
|
encoder,
|
|
num_classes,
|
|
embed_dim,
|
|
representation_type=meta_args.get("representation_type", "last_avgpool"),
|
|
head_type=meta_args.get("head_type", "linear"),
|
|
).to(device)
|
|
|
|
checkpoint_path = _resolve_checkpoint_path(meta_args)
|
|
_load_model_state(model, checkpoint_path, device)
|
|
model.eval()
|
|
|
|
eval_transform = make_transform_eval(crop_size=crop_size)
|
|
split = data_args.get("split", "test")
|
|
root_path = data_args.get("root_path", "./wilds_data")
|
|
|
|
full_dataset, eval_dataset, eval_loader = make_iwildcam_eval(
|
|
transform=eval_transform,
|
|
split=split,
|
|
batch_size=data_args.get("batch_size", 128),
|
|
root_path=root_path,
|
|
num_workers=data_args.get("num_workers", 8),
|
|
pin_mem=bool(data_args.get("pin_mem", True)),
|
|
download=bool(data_args.get("download", True)),
|
|
)
|
|
|
|
all_y_pred = []
|
|
all_y_true = []
|
|
all_metadata = []
|
|
|
|
with torch.no_grad():
|
|
for batch in eval_loader:
|
|
if len(batch) == 3:
|
|
imgs, y_true, metadata = batch
|
|
else:
|
|
raise ValueError("Expected eval loader to return (x, y, metadata)")
|
|
imgs = imgs.to(device, non_blocking=True)
|
|
with torch.cuda.amp.autocast(enabled=use_autocast, dtype=torch.bfloat16):
|
|
logits = model(imgs)
|
|
preds = logits.argmax(dim=1).cpu()
|
|
all_y_pred.append(preds)
|
|
all_y_true.append(y_true.cpu())
|
|
all_metadata.append(metadata.cpu())
|
|
|
|
if all_y_pred:
|
|
all_y_pred = torch.cat(all_y_pred, dim=0)
|
|
all_y_true = torch.cat(all_y_true, dim=0)
|
|
all_metadata = torch.cat(all_metadata, dim=0) if all_metadata else torch.empty(0)
|
|
|
|
if int(all_metadata.shape[0]) != int(all_y_pred.shape[0]):
|
|
raise ValueError(
|
|
"Metadata length mismatch with predictions: "
|
|
f"{int(all_metadata.shape[0])} vs {int(all_y_pred.shape[0])}"
|
|
)
|
|
|
|
metrics = full_dataset.eval(all_y_pred, all_y_true, all_metadata)
|
|
|
|
metrics_path = os.path.join(folder, f"{tag}_metrics.json")
|
|
with open(metrics_path, "w") as f:
|
|
json.dump(metrics, f, indent=2, sort_keys=True)
|
|
logger.info(f"Eval metrics saved to {metrics_path}")
|
|
logger.info(f"Eval metrics: {metrics}")
|
|
|
|
if (
|
|
not force_single
|
|
and dist.is_available()
|
|
and dist.is_initialized()
|
|
and world_size > 1
|
|
):
|
|
dist.barrier()
|
|
dist.destroy_process_group()
|
|
|
|
return {
|
|
"metrics": metrics,
|
|
"metrics_path": metrics_path,
|
|
"folder": folder,
|
|
}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise RuntimeError(
|
|
"Use main_eval_wilds.py to launch this script with a config file."
|
|
)
|