| |
| """Render per-decoder-stage PCA features from a pretraining checkpoint. |
| |
| python scripts/visualize_ckpt.py \ |
| --config configs/aff_base_pretrain_0.4ds_0.5mask_last_local.yaml \ |
| --checkpoint weights/pretrain/ckpt_epoch_399_affmae_fpw.pth \ |
| --samples docs/assets/sample{1,2,3,4}.png \ |
| --output output/plots/pca_decoder_stages.png |
| |
| Without ``--samples`` it pulls a batch from the config's pretraining shards |
| instead, which is what the training loop does; naming files is usually what you |
| want for a figure, since the same images then appear in every run. |
| |
| The model runs **unmasked**: ``forward_without_masking`` queries the decoder at |
| every patch position, so each stage's PCA covers the whole image with no holes. |
| |
| This used to carry its own copy of the renderer, along with copies of |
| ``compute_pca_rgb`` and ``denormalize``, and that copy ran ``_forward_internal`` |
| instead -- so it masked the input and painted PCA only inside the masked region, |
| leaving the visible image showing through everywhere else. Two divergent |
| renderers for one figure, and the masked one is not what you want to look at: |
| the holes are where the model had no input, not where its features are |
| interesting. It now calls the same function the training loop uses. |
| """ |
|
|
| import argparse |
| import os |
| import sys |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
|
|
| import torch |
|
|
| from affmae.config import load_config |
| from affmae.models.registry import get_model_spec |
| from affmae.utils.misc import strip_module_prefix |
| from affmae.viz.model_figures import run_pca_visualization |
|
|
|
|
|
|
| def load_samples(paths, config): |
| """Preprocess image files exactly as the pretraining loader would. |
| |
| Reuses ``apply_custom_processing`` and ``create_transforms`` rather than |
| ``preprocess_image``: the latter normalizes with the *finetuning* statistics |
| and skips the microscope info-bar crop, and both matter here. The bar is a |
| bright strip the encoder never saw during pretraining, because that crop |
| removed it, so leaving it in puts out-of-distribution input at the bottom of |
| every PCA panel. |
| |
| Args: |
| paths: image files, one row of the figure each. |
| config: Config, read for ``img_size`` and ``in_channels``. |
| Returns: |
| [len(paths), C, img_size, img_size] normalized tensor. |
| """ |
| from PIL import Image |
|
|
| from affmae.data.pretrain_dataset import ( |
| apply_custom_processing, |
| create_transforms, |
| ) |
|
|
| transform = create_transforms(config.img_size, config.in_channels) |
| tensors = [] |
| for path in paths: |
| if not os.path.isfile(path): |
| raise SystemExit(f"sample not found: {path}") |
| |
| |
| sample = {"__key__": path, "__url__": "", "png": Image.open(path)} |
| tensors.append(transform(apply_custom_processing(sample)["png"])) |
| return torch.stack(tensors) |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser( |
| description="PCA of each decoder stage's features, unmasked.") |
| parser.add_argument("--config", required=True, |
| help="Pretraining config: model architecture, img_size, " |
| "and data.path when --samples is not given.") |
| parser.add_argument("--samples", nargs="+", default=None, |
| help="Image files to use instead of the shards.") |
| parser.add_argument("--checkpoint", required=True, help="Checkpoint (.pth).") |
| parser.add_argument("--output", required=True, help="Output image (.png/.pdf).") |
| parser.add_argument("--batch-skip", type=int, default=6, |
| help="Batches to discard before visualizing.") |
| parser.add_argument("--num-images", type=int, default=6, |
| help="How many images to visualize.") |
| parser.add_argument("--device", default=None, |
| help="Defaults to cuda when available.") |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
|
|
| device_name = args.device or ("cuda" if torch.cuda.is_available() else "cpu") |
| if not os.path.exists(args.config): |
| raise SystemExit(f"Config file not found: {args.config}") |
|
|
| config = load_config(args.config) |
| config.device = device_name |
| if not args.samples and not getattr(config, "path", None): |
| raise SystemExit( |
| "without --samples this reads the pretraining shards, so it needs a " |
| "pretraining config with data.path set. A finetuning config has " |
| "data.base_path instead and will not work here.") |
|
|
| device = torch.device(device_name) |
| spec = get_model_spec(config.model_type) |
| if spec.build_pretrain is None: |
| raise SystemExit(f"model '{spec.name}' has no pretraining variant.") |
| print(f"Initializing model type: {config.model_type}") |
| model = spec.build_pretrain(config).to(device) |
|
|
| if not os.path.isfile(args.checkpoint): |
| raise SystemExit(f"Checkpoint file not found: {args.checkpoint}") |
| checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False) |
| state_dict = strip_module_prefix( |
| checkpoint.get("model_state_dict", checkpoint)) |
| missing, unexpected = model.load_state_dict(state_dict, strict=True) |
| print(f"Weights loaded. Missing: {len(missing)}, Unexpected: {len(unexpected)}") |
|
|
| if args.samples: |
| batch_images = load_samples(args.samples, config) |
| else: |
| from affmae.data.pretrain_dataset import build_pretrain_dataloader |
|
|
| dataloader, _ = build_pretrain_dataloader(config) |
| data_iter = iter(dataloader) |
| try: |
| for _ in range(args.batch_skip): |
| next(data_iter) |
| batch_images, _ = next(data_iter) |
| except StopIteration: |
| raise SystemExit( |
| f"the loader ran out of data before batch {args.batch_skip}; " |
| f"lower --batch-skip.") |
|
|
| count = min(batch_images.shape[0], args.num_images) |
| print(f"Visualizing {count} images, unmasked...") |
| os.makedirs(os.path.dirname(os.path.abspath(args.output)) or ".", exist_ok=True) |
| run_pca_visualization(model, batch_images[:count].to(device), args.output, |
| device) |
| print(f"Wrote {args.output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
