File size: 4,560 Bytes
93ffd19 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | #!/usr/bin/env python
from __future__ import annotations
import argparse
from pathlib import Path
import copy
import torch
import torch.nn.functional as F
from tqdm import tqdm
from sacflow.utils.config import load_yaml
from sacflow.utils.misc import seed_everything, ensure_dir, move_to_device
from sacflow.utils.distributed import init_distributed, cleanup, is_main_process
from sacflow.data.loader import build_loader
from sacflow.models.unet3d import build_model
from sacflow.methods.task_space import centered_classifier_basis, project_task_and_residual
from sacflow.methods.source_memory import hard_onehot, class_moments, save_source_memory
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", required=True)
ap.add_argument("--checkpoint", required=True)
ap.add_argument("--output", required=True)
ap.add_argument("--num-passes", type=int, default=3, help="Number of random-crop passes over source_train for memory estimation")
args = ap.parse_args()
cfg = load_yaml(args.config)
# Source memory should estimate source feature statistics, not augmentation noise.
cfg = copy.deepcopy(cfg)
cfg.setdefault("data", {}).setdefault("augmentation", {})
cfg["data"]["augmentation"] = {"random_flip": False, "random_intensity_shift": 0.0, "random_intensity_scale": 0.0}
seed_everything(int(cfg.get("seed", 1337)))
device = init_distributed(cfg.get("distributed", {}).get("backend", "nccl"))
model = build_model(cfg).to(device)
ckpt = torch.load(args.checkpoint, map_location="cpu")
model.load_state_dict(ckpt.get("model", ckpt), strict=False)
model.eval()
loader = build_loader(cfg, split="source_train", training=True, require_label=True)
W = model.final_classifier_weight().detach().to(device)
Q = centered_classifier_basis(W)
C = cfg["data"]["num_classes"]
mu_acc = None
var_acc = None
feat_mu_acc = None
feat_var_acc = None
count_acc = None
with torch.no_grad():
for pass_idx in range(max(1, args.num_passes)):
for batch in tqdm(loader, desc=f"source memory pass {pass_idx+1}/{max(1,args.num_passes)}", disable=not is_main_process()):
batch = move_to_device(batch, device)
logits, feats = model(batch["image"], return_features=True)
feat = feats["prelogit"]
task, residual = project_task_and_residual(feat, Q)
labels = batch["label"]
if labels.shape[-3:] != residual.shape[-3:]:
labels = F.interpolate(labels[:,None].float(), size=residual.shape[-3:], mode="nearest")[:,0].long()
probs = hard_onehot(labels, C)
mu, std, counts = class_moments(residual, probs)
fmu, fstd, _ = class_moments(feat, probs)
var = std.pow(2)
fvar = fstd.pow(2)
if mu_acc is None:
mu_acc = mu * counts[:,None]
var_acc = (var + mu.pow(2)) * counts[:,None]
feat_mu_acc = fmu * counts[:,None]
feat_var_acc = (fvar + fmu.pow(2)) * counts[:,None]
count_acc = counts
else:
mu_acc += mu * counts[:,None]
var_acc += (var + mu.pow(2)) * counts[:,None]
feat_mu_acc += fmu * counts[:,None]
feat_var_acc += (fvar + fmu.pow(2)) * counts[:,None]
count_acc += counts
count = count_acc.clamp_min(1e-6)
mu = mu_acc / count[:,None]
second = var_acc / count[:,None]
var = (second - mu.pow(2)).clamp_min(1e-6)
fmu = feat_mu_acc / count[:,None]
fsecond = feat_var_acc / count[:,None]
fvar = (fsecond - fmu.pow(2)).clamp_min(1e-6)
mem = {
"classifier_weight": W.detach().cpu(),
"task_basis_Q": Q.detach().cpu(),
"residual_mu": mu.detach().cpu(),
"residual_std": torch.sqrt(var).detach().cpu(),
"feature_mu": fmu.detach().cpu(),
"feature_std": torch.sqrt(fvar).detach().cpu(),
"class_counts": count.detach().cpu(),
"num_classes": C,
"feature_dim": W.shape[1],
"note": "Compact source memory. No raw source images stored. Estimated from source_train random crops with augmentation disabled.",
"num_passes": args.num_passes,
}
if is_main_process():
save_source_memory(args.output, mem)
print(f"Saved source memory to {args.output}")
cleanup()
if __name__ == "__main__":
main()
|