#!/usr/bin/env python """Cache SAE activations for analysis notebooks and dashboard. Usage: python analysis/tools/cache_activations.py python analysis/tools/cache_activations.py --checkpoint /path/to/checkpoint python analysis/tools/cache_activations.py --max-batches 10 python analysis/tools/cache_activations.py --max-memory-gb 8 Paths and defaults are read from ``default_configs/sae_vis_config.json`` (or ``SAE_VIS_CONFIG_PATH`` / ``--sae-vis-config-path``). """ from __future__ import annotations import argparse import json import os import sys from pathlib import Path ROOT = Path(__file__).resolve().parents[2] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from analysis.cache_utils import build_activation_cache from analysis.config import load_sae_vis_config from log_config import get_logger, setup_logging logger = get_logger(__name__) def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description='Cache KADID-10k SAE activations') p.add_argument( '--sae-vis-config-path', type=str, default=None, help='Path to sae_vis_config.json (default: template or SAE_VIS_CONFIG_PATH)', ) p.add_argument( '--dataset', type=str, default=None, choices=['kadid', 'local_kadid'], help='Explicit dataset choice for caching', ) p.add_argument('--checkpoint', type=str, default=None, help='Override SAE_CHECKPOINT_PATH') p.add_argument('--kadid-path', type=str, default=None, help='Override KADID_IMAGES_PATH') p.add_argument('--layer', type=int, default=None, help='Override LAYER_NUM') p.add_argument( '--iqa-metric', type=str, default=None, choices=['arniqa-kadid', 'maniqa', 'qalign', 'liqe', 'liqe_mix', 'topiq_nr'], help='Override IQA_METRIC', ) p.add_argument( '--swin-num', type=int, default=None, choices=[1, 2], help='Override SWIN_NUM (MANIQA only)', ) p.add_argument('--device', type=str, default=None, help='Override DEVICE (cuda / cpu)') p.add_argument('--batch-size', type=int, default=None, help='Override BATCH_SIZE') p.add_argument('--num-workers', type=int, default=None, help='Override NUM_WORKERS') p.add_argument('--max-batches', type=int, default=None, help='Limit batches (debug)') p.add_argument( '--max-memory-gb', type=float, default=30.0, help='Memory limit for accumulated sparse matrices, GB', ) p.add_argument( '--min-distortion-level', type=int, default=None, help='Minimum KADID distortion level (1..5)', ) p.add_argument( '--srground-include-sr-artifact', action='store_true', help='Include SR artifacts in SRGround mask (default: real distortions only)', ) return p.parse_args() def main() -> None: setup_logging() args = parse_args() cfg = load_sae_vis_config(args.sae_vis_config_path) dataset_arg = args.dataset if dataset_arg is None: dataset_arg = 'local_kadid' if cfg.DATASET == 'local_kadid' else 'kadid' dataset_key = 'local_kadid' if dataset_arg == 'local_kadid' else 'kadid10k' checkpoint_path = args.checkpoint or cfg.SAE_CHECKPOINT_PATH dataset_root = args.kadid_path or cfg.DATASET_ROOT layer_num = args.layer if args.layer is not None else cfg.LAYER_NUM iqa_metric = args.iqa_metric or cfg.IQA_METRIC swin_num = args.swin_num if args.swin_num is not None else cfg.SWIN_NUM device = args.device or cfg.DEVICE batch_size = args.batch_size if args.batch_size is not None else cfg.BATCH_SIZE num_workers = args.num_workers if args.num_workers is not None else cfg.NUM_WORKERS min_distortion_level = ( args.min_distortion_level if args.min_distortion_level is not None else cfg.KADID_MIN_DISTORTION_LEVEL ) if not (1 <= min_distortion_level <= 5): raise ValueError('--min-distortion-level must be in [1, 5]') if args.checkpoint is not None: cache_dir = os.path.join(os.path.dirname(checkpoint_path), 'activation_cache') acts_filename = os.path.basename(cfg.DATASET_CACHE_CONFIGS[dataset_key].acts_cache_path) cache_path = os.path.join(cache_dir, dataset_key, acts_filename) os.makedirs(os.path.dirname(cache_path), exist_ok=True) else: cache_path = cfg.DATASET_CACHE_CONFIGS[dataset_key].acts_cache_path logger.info('=' * 60) logger.info(' SAE checkpoint : %s', checkpoint_path) logger.info(' Dataset : %s', dataset_arg) logger.info(' Dataset root : %s', dataset_root) logger.info(' Layer : %s', layer_num) logger.info(' IQA metric : %s', iqa_metric) if iqa_metric == 'maniqa': logger.info(' Swin num : %s', swin_num) logger.info(' Device : %s', device) logger.info(' Batch size : %s', batch_size) logger.info(' Min dist level : %s', min_distortion_level) logger.info(' Cache output : %s', cache_path) if args.max_batches: logger.debug(' Max batches : %s', args.max_batches) logger.info('=' * 60) logger.info('Starting caching...') build_info = build_activation_cache( dataset=dataset_key, cache_path=cache_path, checkpoint_path=checkpoint_path, dataset_root=dataset_root, layer_num=layer_num, iqa_metric=iqa_metric, swin_num=swin_num, device=device, batch_size=batch_size, num_workers=num_workers, crop_size=cfg.CROP_SIZE, scaling_factor=cfg.SCALING_FACTOR, min_distortion_level=min_distortion_level, max_batches=args.max_batches, max_memory_gb=args.max_memory_gb, add_patch_mask_stats=True, include_pristine=True, show_progress_bars=cfg.SHOW_PROGRESS_BARS, srground_include_sr_artifact=( args.srground_include_sr_artifact or cfg.SRGROUND_INCLUDE_SR_ARTIFACT ), ) layer_name = str(build_info['layer_name']) patch_grid_shape = tuple(build_info['patch_grid_shape']) patches_per_image = int(build_info['patches_per_image']) sae_config = build_info.get('sae_config') logger.info(' Hooked layer: %r', layer_name) logger.info(' Grid=%s, patches/image=%s', patch_grid_shape, patches_per_image) if sae_config: logger.info(' SAE type: %s', sae_config.get('sae_type', 'unknown')) logger.info(' Lambda param: %s', sae_config.get('lambda_param', 'unknown')) logger.info(' Inner dim: %s', sae_config.get('inner_dim', 'unknown')) run_meta = { 'dataset_type': dataset_arg, 'iqa_metric': iqa_metric, 'swin_num': swin_num, 'layer': layer_num, 'batch_size': batch_size, 'crop_size': cfg.CROP_SIZE, 'min_distortion_level': min_distortion_level, 'checkpoint_path': checkpoint_path, 'cache_path': cache_path, } run_meta_path = os.path.splitext(cache_path)[0] + '_run_meta.json' with open(run_meta_path, 'w', encoding='utf-8') as f: json.dump(run_meta, f, indent=2, ensure_ascii=True) logger.info('Run metadata saved -> %s', run_meta_path) logger.info('Done.') if __name__ == '__main__': main()