Zennit Crp

Zennit-CRP: Concept Relevance Propagation Skill

When to Use

Activate this skill when:

• Generating explainability heatmaps for deep learning model predictions
• Computing concept-conditional attributions (CRP) to understand which latent features drive predictions
• Performing Relevance Maximization (RelMax) or Activation Maximization (ActMax) to find representative samples for neuron/channel concepts
• Building attribution graphs to trace relevance flows through network layers
• Analyzing channel-level importance in convolutional or linear layers
• Comparing what concepts a model uses vs. what activates a neuron (relevance vs. activation)
• Working with LRP (Layer-wise Relevance Propagation) via zennit composites
• Any task involving keywords: CRP, LRP, XAI, explainable AI, feature visualization, relevance propagation, concept attribution, channel relevance, RelMax, ActMax

Quick Reference

• Paper: Nature Machine Intelligence (Open Access)
• PyPI: https://pypi.org/project/zennit-crp/
• GitHub: https://github.com/rachtibat/zennit-crp
• Tutorials: https://github.com/rachtibat/zennit-crp/tree/master/tutorials
• Zennit (dependency): https://github.com/chr5tphr/zennit
• Citation: Achtibat et al., "From attribution maps to human-understandable explanations through Concept Relevance Propagation", Nature Machine Intelligence, 2023

Installation / Setup

Prerequisites

• Python 3.7+
• PyTorch (with CUDA support recommended for large datasets)
• A working model in PyTorch (torch.nn.Module)

Install from PyPI (recommended, includes fast image utilities)

pip install zennit-crp[fast_img]

Install from Source (to access tutorials)

git clone https://github.com/rachtibat/zennit-crp
pip install ./zennit-crp

Install zennit separately if needed

pip install zennit

Core Features

• Conditional Attributions (CRP): Generate concept-conditional relevance heatmaps by masking relevance flows during backpropagation, isolating the contribution of individual latent concepts to the model output.
• ChannelConcept: Define each channel in Conv2D/Linear layers as a distinct concept with built-in relative importance scoring.
• Custom Concepts: Extend the abstract Concept class to define your own notion of what constitutes a concept (beyond channels).
• Relevance Maximization (RelMax): Select dataset samples that maximize the relevance criterion for a given neuron/channel — a more faithful alternative to Activation Maximization (ActMax).
• Activation Maximization (ActMax): Classic approach: select samples that most strongly activate a given unit.
• Feature Visualization: The FeatureVisualization class precomputes and caches reference images (RelMax/ActMax) for all concepts across all layers on a full dataset.
• Attribution Graph: Trace relevance as a directed graph through model layers to understand information flow from output back to input.
• Model Graph Analysis: Automatically parse PyTorch JIT-traced model graphs to determine layer connectivity for attribution routing.
• LRP Integration: Built on top of zennit composites (e.g., EpsilonPlusFlat, EpsilonAlpha2Beta1) and canonizers (e.g., SequentialMergeBatchNorm).
• Caching System: ImageCache and abstract Cache classes for efficient storage and retrieval of reference images.
• Image Utilities: plot_grid and related helpers in crp.image for visualizing attribution results and reference image grids.
• Statistics Module: Tools in crp.statistics for aggregating concept relevance statistics across datasets.

Usage Examples

Conditional Attributions (CRP Quickstart)

from crp.attribution import CondAttribution
from crp.concepts import ChannelConcept
from crp.helper import get_layer_names

from zennit.composites import EpsilonPlusFlat
from zennit.canonizers import SequentialMergeBatchNorm
import torch.nn as nn

# define LRP rules and canonizers in zennit
composite = EpsilonPlusFlat([SequentialMergeBatchNorm()])

# load CRP toolbox
attribution = CondAttribution(model)

# here, each channel is defined as a concept
# or define your own notion!
cc = ChannelConcept()

# get layer names of Conv2D and MLP layers
layer_names = get_layer_names(model, [nn.Conv2d, nn.Linear])

# get a conditional attribution for channel 50 in layer features.27 wrt. output 1
conditions = [{'features.27': [50], 'y': [1]}]

attr = attribution(data, conditions, composite, record_layer=layer_names)

# heatmap and prediction
attr.heatmap, attr.prediction
# activations and relevances for each layer name
attr.activations, attr.relevances

# relative importance of each concept for final prediction
rel_c = cc.attribute(attr.relevances['features.40'])
# most relevant channels in features.40
concept_ids = torch.argsort(rel_c, descending=True)

Feature Visualization (RelMax / ActMax)

from crp.visualization import FeatureVisualization
from crp.image import plot_grid

# define which concept is used in each layer
layer_map = {name: cc for name in layer_names}

# compute visualization (it takes for VGG16 and ImageNet testset on Titan RTX ~30 min)
fv = FeatureVisualization(attribution, dataset, layer_map)
fv.run(composite, 0, len(dataset))

# visualize MaxRelevance reference images for top-5 concepts
ref_c = fv.get_max_reference(concept_ids[:5], 'features.40', 'relevance', composite=composite)

plot_grid(ref_c)

Key APIs / Models

Classes

Key Functions

Zennit Composites (used with CRP)

• EpsilonPlusFlat — LRP-ε for middle layers, LRP-0+ for early layers
• EpsilonAlpha2Beta1 — LRP with alpha-beta rules
• GuidedBackprop — Guided backpropagation composite

Zennit Canonizers

• SequentialMergeBatchNorm — Folds BatchNorm into preceding Conv/Linear layers

Condition Dictionary Format

# Condition on specific channels and output class
conditions = [{'layer_name': [channel_ids], 'y': [class_ids]}]

# Unconditional (standard LRP)
conditions = [{'y': [1]}]

# Multi-layer condition
conditions = [{'features.27': [50, 100], 'features.15': [32], 'y': [1]}]

Common Patterns & Best Practices

1. Always merge BatchNorm before attribution

When your model uses BatchNorm layers, include SequentialMergeBatchNorm() in the canonizers list to ensure numerically stable LRP:

composite = EpsilonPlusFlat([SequentialMergeBatchNorm()])

2. Record intermediate activations/relevances

Pass record_layer=layer_names to attribution() to capture activations and relevances at every layer for downstream analysis:

attr = attribution(data, conditions, composite, record_layer=layer_names)
# Access per-layer data:
print(attr.relevances['features.27'].shape)
print(attr.activations['features.27'].shape)

3. Sort concepts by relevance for top-k analysis

rel_c = cc.attribute(attr.relevances['features.40'])
top_k_concepts = torch.argsort(rel_c, descending=True)[:10]

4. Run FeatureVisualization once, query many times

fv.run(...) is expensive (precomputes over entire dataset). Run it once and save results. Then use fv.get_max_reference(...) for fast retrieval.

5. Use 'relevance' mode over 'activation' for faithful explanations

RelMax selects samples based on how useful a concept is for the model's classification task, making it more faithful than activation-based selection in adversarial scenarios.

6. Layer naming convention

Layer names follow PyTorch's named_modules() format (e.g., 'features.27', 'classifier.6'). Use get_layer_names(model, [nn.Conv2d, nn.Linear]) to enumerate valid names automatically.

7. Batched conditions

CondAttribution supports a list of condition dicts (one per sample in the batch), enabling efficient batched conditional attribution.

conditions = [
    {'features.27': [50], 'y': [1]},
    {'features.27': [100], 'y': [3]},
]
attr = attribution(batch_data, conditions, composite, record_layer=layer_names)

Reference-Image Cropping for Foundation-Model Embedding

fv.get_max_reference(...) returns whole input images. When the downstream goal is to embed each reference image into a vision foundation model (CLIP, SigLIP, DINOv2, ...) so that every neuron gets a concept vector in that model's semantic space, the full image is usually too noisy — the foundation model embeds the background, not the concept. The standard fix is to crop each reference image to the high-relevance region using the CRP heatmap itself. The recipe is identical across encoders; only the Gaussian blur kernel scales with input resolution.

Standard pipeline

import torch
from torchvision.transforms.functional import gaussian_blur
from crp.image import get_crop_range

def _square_box(r1, r2, c1, c2):
    """Symmetric expansion of a tight bbox into a square box (no inflate)."""
    dr, dc = r2 - r1, c2 - c1
    if dr > dc:
        pad = (dr - dc) // 2
        c1 -= pad; c2 += pad
        if c1 < 0:
            c2 -= c1; c1 = 0
    elif dc > dr:
        pad = (dc - dr) // 2
        r1 -= pad; r2 += pad
        if r1 < 0:
            r2 -= r1; r1 = 0
    return r1, r2, c1, c2

def crop_to_concept(image, heatmap, *, kernel_size, crop_th, square=True):
    """Blur the heatmap, threshold at `crop_th * max`, return the tight bbox.

    image:    (C, H, W) input tensor in the same resolution as the heatmap.
    heatmap:  (H, W)    per-pixel relevance from CondAttribution.
    """
    h = gaussian_blur(heatmap.unsqueeze(0), kernel_size=kernel_size)[0]
    h = h.abs() / (h.abs().max() + 1e-8)
    r1, r2, c1, c2 = get_crop_range(h, crop_th)
    if square:
        r1, r2, c1, c2 = _square_box(r1, r2, c1, c2)
    return image[..., r1:r2, c1:c2]

Pattern: blur → abs → normalize-by-max → threshold → tight bbox → (optional) squarify.

• Blur prevents 1-2 noisy hot pixels at the image border from anchoring the bbox.
• abs + normalize turns signed LRP relevance into a [0, 1] saliency.
• Threshold at crop_th × max keeps every pixel above a fraction of peak relevance — crop_th is the only knob controlling crop tightness.
• Squarify avoids non-square crops being resize-distorted by the downstream encoder; expansion is symmetric and stays inside the image, so there is no padding around the heatmap region.

Presets

zennit-crp's vis_opaque_img and vis_img_heatmap ship with defaults tuned for human heatmap inspection — small kernel, strict thresholds, rectangular box. Embedding crops into a vision foundation model wants the opposite: bigger blur (smoother bbox boundary), 10× more permissive thresholds (more context for the encoder), forced square box (no resize distortion).

Scaling rule for a non-standard encoder side H: kernel ≈ 0.23 × H, rounded to the nearest odd integer (torchvision.transforms.functional.gaussian_blur requires odd kernels). For DINOv2-base at 224 or SigLIP-base at 256, plug H into the same rule rather than reading the table literally.

The three sub-1 thresholds are encoder-invariant: 1 % of max relevance keeps essentially all of the concept's mass while excluding background, which is a property of natural-image LRP/CRP heatmaps, not of any specific encoder. The same crop_th / vis_th / alpha values transfer across CLIP, SigLIP, DINOv2, and any other natural-image foundation model.

Picking a preset

• Building a UMAP / cluster view for a person to look at → human-inspection defaults. Tight crops are fine because the human eye fills in context, and you usually want to see exactly where on the image the network looks.
• CLIP / OpenCLIP / MobileCLIP / EVA-CLIP for image-text concept labelling → 224 preset. Default choice when each neuron should be labelled with a text concept via cosine similarity in CLIP space.
• SigLIP / SigLIP-2 for higher-resolution image-text embedding → 384 preset. Use when CLIP-224 patches lose too much detail — fine-grained categories, small objects, dense scenes, medical or satellite imagery.
• DINOv2 (or any vision-only foundation model) for clustering / similarity → high-res preset. Use when you want a semantic embedding without the text-alignment bias that CLIP-family models inherit from contrastive pre-training (e.g., comparing what two networks visually group together, independent of how that grouping maps to language).

Why the embedding thresholds are 10× more permissive than human-inspection defaults

vis_opaque_img's crop_th=0.1 keeps only pixels above 10 % of max relevance — typically just the brightest concept core. A person looking at the heatmap fills in surrounding context from memory. A foundation encoder cannot: cropped that tightly, CLIP / SigLIP / DINOv2 see a meaningless texture patch and embed noise. Dropping crop_th to 0.01 recovers everything at ≥1 % of max relevance, making the crop large enough for the encoder to actually recognize what concept is there.

CRP heatmap source

For CNN channels, CondAttribution returns an input-space heatmap directly. For Vision Transformers, CRP through self-attention is not implemented in zennit-crp at the time of writing; substitute the heatmap with the upsampled spatial-token activation map (bilinear-interpolated to image resolution) and feed it through crop_to_concept as-is. Treat ViT crops as a fallback — they are an activation-map proxy, not true relevance.

Demo Scripts

scripts/conditional_attribution_demo.py

#!/usr/bin/env python3
"""
Conditional Attribution Demo using zennit-crp (Concept Relevance Propagation)

This script demonstrates how to:
  1. Set up a PyTorch model with zennit LRP composites
  2. Compute concept-conditional relevance heatmaps using CRP
  3. Identify the most relevant channels/concepts in a given layer
  4. Record per-layer activations and relevances

Requires:
    pip install zennit-crp[fast_img] torch torchvision

Usage:
    python conditional_attribution_demo.py
"""

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import numpy as np

# CRP imports
from crp.attribution import CondAttribution
from crp.concepts import ChannelConcept
from crp.helper import get_layer_names

# Zennit LRP composites and canonizers
from zennit.composites import EpsilonPlusFlat
from zennit.canonizers import SequentialMergeBatchNorm


# ---------------------------------------------------------------------------
# Model Setup
# ---------------------------------------------------------------------------

def load_model() -> nn.Module:
    """
    Load a pretrained VGG16 model from torchvision.

    Returns:
        model (nn.Module): A pretrained VGG16 in eval mode.
    """
    model = models.vgg16(pretrained=True)
    model.eval()
    return model


def create_dummy_input(batch_size: int = 2) -> torch.Tensor:
    """
    Create a random image batch as a stand-in for real data.
    Replace this with actual image loading for real use.

    Args:
        batch_size: Number of images in the batch.

    Returns:
        Tensor of shape (batch_size, 3, 224, 224) with requires_grad=True.
    """
    # In real usage, load actual images and apply standard ImageNet preprocessing:
    # transform = transforms.Compose([
    #     transforms.Resize(256),
    #     transforms.CenterCrop(224),
    #     transforms.ToTensor(),
    #     transforms.Normalize(mean=[0.485, 0.456, 0.406],
    #                          std=[0.229, 0.224, 0.225]),
    # ])
    # img = Image.open("path/to/image.jpg").convert("RGB")
    # data = transform(img).unsqueeze(0)

    data = torch.randn(batch_size, 3, 224, 224, requires_grad=True)
    return data


# ---------------------------------------------------------------------------
# LRP Setup
# ---------------------------------------------------------------------------

def build_composite() -> EpsilonPlusFlat:
    """
    Build a zennit EpsilonPlusFlat composite with BatchNorm canonizer.
    EpsilonPlusFlat uses:
        - LRP-epsilon for middle layers (numerically stable)
        - LRP-0+ (flat) for the input/first layers (avoids gradient shattering)

    Returns:
        composite (EpsilonPlusFlat): Ready-to-use LRP composite.
    """
    canonizers = [SequentialMergeBatchNorm()]
    composite = EpsilonPlusFlat(canonizers)
    return composite


# ---------------------------------------------------------------------------
# CRP Attribution
# ---------------------------------------------------------------------------

def run_conditional_attribution(
    model: nn.Module,
    data: torch.Tensor,
    target_layer: str = "features.28",
    concept_ids: list = None,
    output_class: int = 1,
) -> dict:
    """
    Run Concept Relevance Propagation (CRP) for specified concept IDs
    in a given layer, conditioned on a target output class.

    Args:
        model:        The PyTorch model (eval mode).
        data:         Input tensor, shape (B, C, H, W), requires_grad=True.
        target_layer: Name of the layer to condition on (e.g., 'features.28').
        concept_ids:  List of channel indices to condition on. Defaults to [50].
        output_class: The output class index to attribute from.

    Returns:
        Dictionary containing:
            'heatmap'     - Input-space relevance heatmap tensor
            'prediction'  - Model output logits
            'relevances'  - Dict[layer_name -> relevance tensor]
            'activations' - Dict[layer_name -> activation tensor]
            'rel_c'       - Per-channel relative relevance in target_layer
            'top_concepts'- Sorted channel indices by descending relevance
    """
    if concept_ids is None:
        concept_ids = [50]

    # --- Build composite and concept object ---
    composite = build_composite()
    cc = ChannelConcept()

    # --- Discover layer names for Conv2d and Linear layers ---
    layer_names = get_layer_names(model, [nn.Conv2d, nn.Linear])
    print(f"[INFO] Found {len(layer_names)} Conv2d/Linear layers.")
    print(f"[INFO] First 5 layer names: {layer_names[:5]}")

    # --- Initialize CondAttribution ---
    attribution = CondAttribution(model)

    # --- Define conditions ---
    # Each entry in the list corresponds to one sample in the batch.
    # 'y' specifies the output neuron to propagate from.
    # The target_layer key specifies which channel(s) to condition on.
    conditions = [
        {target_layer: concept_ids, "y": [output_class]}
        for _ in range(data.shape[0])
    ]
    print(f"[INFO] Running CRP with conditions: {conditions[0]}")

    # --- Run attribution ---
    # record_layer captures activations AND relevances at each named layer
    attr = attribution(
        data,
        conditions,
        composite,
        record_layer=layer_names,
    )

    # --- Extract results ---
    heatmap = attr.heatmap       # shape: (B, H, W) - input-space relevance
    prediction = attr.prediction  # shape: (B, num_classes) - raw logits

    print(f"[INFO] Heatmap shape:    {heatmap.shape}")
    print(f"[INFO] Prediction shape: {prediction.shape}")

    # --- Compute per-channel relative relevance in the last feature layer ---
    last_layer = layer_names[-1]
    if last_layer in attr.relevances:
        rel_c = cc.attribute(attr.relevances[last_layer])
        top_concepts = torch.argsort(rel_c, descending=True)
        print(f"[INFO] Top-5 concept IDs in '{last_layer}': {top_concepts[:5].tolist()}")
    else:
        rel_c = None
        top_concepts = None
        print(f"[WARN] Layer '{last_layer}' not found in recorded relevances.")

    return {
        "heatmap": heatmap,
        "prediction": prediction,
        "relevances": attr.relevances,
        "activations": attr.activations,
        "rel_c": rel_c,
        "top_concepts": top_concepts,
        "layer_names": layer_names,
    }


# ---------------------------------------------------------------------------
# Unconditional Attribution (standard LRP)
# ---------------------------------------------------------------------------

def run_unconditional_attribution(
    model: nn.Module,
    data: torch.Tensor,
    output_class: int = 1,
) -> object:
    """
    Run standard (unconditional) LRP attribution without concept masking.

    Args:
        model:        The PyTorch model.
        data:         Input tensor.
        output_class: Output class to propagate from.

    Returns:
        Attribution result object with .heatmap and .prediction.
    """
    composite = build_composite()
    attribution = CondAttribution(model)
    layer_names = get_layer_names(model, [nn.Conv2d, nn.Linear])

    # No layer conditions — only specify output class
    conditions = [{"y": [output_class]} for _ in range(data.shape[0])]

    attr = attribution(data, conditions, composite, record_layer=layer_names)
    print(f"[INFO] Unconditional heatmap shape: {attr.heatmap.shape}")
    return attr


# ---------------------------------------------------------------------------
# Channel Concept: Relative Importance
# ---------------------------------------------------------------------------

def analyze_concept_importance(
    relevances_dict: dict,
    layer_name: str,
) -> torch.Tensor:
    """
    Compute the relative importance of each channel (concept) in a layer.

    Args:
        relevances_dict: Dict mapping layer_name -> relevance tensor from attr.relevances.
        layer_name:      The layer to analyze.

    Returns:
        rel_c (Tensor): Per-channel relative relevance scores, shape (num_channels,).
    """
    cc = ChannelConcept()

    if layer_name not in relevances_dict:
        available = list(relevances_dict.keys())
        raise KeyError(
            f"Layer '{layer_name}' not in recorded relevances. "
            f"Available layers (first 5): {available[:5]}"
        )

    layer_relevance = relevances_dict[layer_name]
    print(f"[INFO] Layer '{layer_name}' relevance tensor shape: {layer_relevance.shape}")

    # cc.attribute collapses spatial dims and normalizes across channels
    rel_c = cc.attribute(layer_relevance)
    print(f"[INFO] rel_c shape: {rel_c.shape}")
    print(f"[INFO] Top-5 concept indices: {torch.argsort(rel_c, descending=True)[:5].tolist()}")
    print(f"[INFO] Top-5 concept relevances: {torch.sort(rel_c, descending=True).values[:5].tolist()}")

    return rel_c


# ---------------------------------------------------------------------------
# Main Demo
# ---------------------------------------------------------------------------

def main():
    """
    Main demonstration of CRP conditional attribution workflow.
    """
    print("=" * 60)
    print("  Zennit-CRP: Conditional Attribution Demo")
    print("=" * 60)

    # 1. Load model
    print("\n[STEP 1] Loading VGG16 model...")
    model = load_model()

    # 2. Create dummy input
    print("\n[STEP 2] Creating dummy input batch (2 images, 224x224)...")
    print("         Replace with real image data for actual use.")
    data = create_dummy_input(batch_size=2)

    # 3. Inspect available layer names
    print("\n[STEP 3] Discovering layer names...")
    layer_names = get_layer_names(model, [nn.Conv2d, nn.Linear])
    print(f"         Total layers: {len(layer_names)}")
    print(f"         All layers: {layer_names}")

    # Pick a target layer for conditioning (second-to-last conv layer in VGG16)
    target_layer = "features.28"  # Conv2d in VGG16 block 5
    print(f"\n         Target layer for CRP conditioning: '{target_layer}'")

    # 4. Run conditional attribution
    print("\n[STEP 4] Running conditional CRP attribution...")
    results = run_conditional_attribution(
        model=model,
        data=data,
        target_layer=target_layer,
        concept_ids=[50, 100],  # condition on channels 50 and 100
        output_class=1,
    )

    # 5. Analyze concept importance in another layer
    print("\n[STEP 5] Analyzing per-channel importance in 'features.26'...")
    try:
        rel_c = analyze_concept_importance(results["relevances"], "features.26")
        top5_ids = torch.argsort(rel_c, descending=True)[:5].tolist()
        print(f"         Most relevant concept IDs: {top5_ids}")
    except KeyError as e:
        print(f"         [WARN] {e}")

    # 6. Run unconditional LRP for comparison
    print("\n[STEP 6] Running unconditional LRP (no concept masking)...")
    attr_unconditional = run_unconditional_attribution(model, data, output_class=1)

    # 7. Summary
    print("\n" + "=" * 60)
    print("  Summary")
    print("=" * 60)
    print(f"  Heatmap shape:       {results['heatmap'].shape}")
    print(f"  Prediction shape:    {results['prediction'].shape}")
    print(f"  Layers recorded:     {len(results['relevances'])}")
    if results["top_concepts"] is not None:
        print(f"  Top concept IDs:     {results['top_concepts'][:5].tolist()}")
    print(f"  Unconditional hmap:  {attr_unconditional.heatmap.shape}")
    print("\n[DONE] Demo complete.")
    print("       Next step: see feature_visualization_demo.py for RelMax/ActMax.")


if __name__ == "__main__":
    main()

scripts/feature_visualization_demo.py

#!/usr/bin/env python3
"""
Feature Visualization Demo using zennit-crp (RelMax / ActMax)

This script demonstrates how to:
  1. Build a FeatureVisualization object over a dataset
  2. Run precomputation of RelMax and ActMax reference images
  3. Query top-k reference images for specific concept IDs
  4. Visualize concept reference images with plot_grid

Requires:
    pip install zennit-crp[fast_img] torch torchvision

Usage:
    python feature_visualization_demo.py

Note:
    Full precomputation on ImageNet + VGG16 takes ~30 min on a GPU (Titan RTX).
    This demo uses a tiny synthetic dataset for illustration.
    Replace `SyntheticDataset` with your real `torch.utils.data.Dataset`.
"""

import torch
import torch.nn as nn
import torchvision.models as models
from torch.utils.data import Dataset, DataLoader
import numpy as np
from pathlib import Path

# CRP imports
from crp.attribution import CondAttribution
from crp.concepts import ChannelConcept
from crp.helper import get_layer_names
from crp.visualization import FeatureVisualization
from crp.image import plot_grid

# Zennit imports
from zennit.composites import EpsilonPlusFlat
from zennit.canonizers import SequentialMergeBatchNorm


# ---------------------------------------------------------------------------
# Synthetic Dataset (replace with your real dataset)
# ---------------------------------------------------------------------------

class SyntheticDataset(Dataset):
    """
    A minimal synthetic dataset for demonstration purposes.

    Replace this with your actual dataset, e.g.:
        from torchvision.datasets import ImageNet
        dataset = ImageNet(root='/path/to/imagenet', split='val',
                           transform=transform)

    The dataset must return (image_tensor, label) tuples,
    where image_tensor has shape (C, H, W).
    """

    def __init__(self, num_samples: int = 50, num_classes: int = 10):
        """
        Args:
            num_samples: Total number of synthetic images.
            num_classes: Number of output classes.
        """
        self.num_samples = num_samples
        self.num_classes = num_classes
        # Simulate ImageNet-like 224x224 RGB images
        self.data = torch.randn(num_samples, 3, 224, 224)
        self.labels = torch.randint(0, num_classes, (num_samples,))

    def __len__(self) -> int:
        return self.num_samples

    def __getitem__(self, idx: int):
        return self.data[idx], self.labels[idx]


# ---------------------------------------------------------------------------
# Model and Composite Setup
# ---------------------------------------------------------------------------

def load_model_and_composite():
    """
    Load VGG16 and build the LRP composite.

    Returns:
        model (nn.Module): VGG16 in eval mode.
        composite (EpsilonPlusFlat): LRP composite with BN canonizer.
    """
    model = models.vgg16(pretrained=True)
    model.eval()
    composite = EpsilonPlusFlat([SequentialMergeBatchNorm()])
    return model, composite


# ---------------------------------------------------------------------------
# Feature Visualization
# ---------------------------------------------------------------------------

def build_feature_visualization(
    model: nn.Module,
    dataset: Dataset,
    layer_names: list,
    save_path: str = "./fv_cache",
) -> FeatureVisualization:
    """
    Construct a FeatureVisualization object for the given model and dataset.

    Args:
        model:       The PyTorch model in eval mode.
        dataset:     A torch Dataset returning (image, label) tuples.
        layer_names: List of layer name strings to compute visualizations for.
        save_path:   Directory path where precomputed results are cached to disk.

    Returns:
        fv (FeatureVisualization): Configured feature visualization object.
    """
    attribution = CondAttribution(model)
    cc = ChannelConcept()

    # Map each layer to the ChannelConcept (one concept per channel)
    layer_map = {name: cc for name in layer_names}

    fv = FeatureVisualization(
        attribution=attribution,
        dataset=dataset,
        layer_map=layer_map,
        path=save_path,  # Cache directory — will be created if it doesn't exist
    )
    return fv


def run_precomputation(
    fv: FeatureVisualization,
    composite,
    start: int = 0,
    stop: int = None,
    batch_size: int = 8,
) -> None:
    """
    Run the precomputation of RelMax and ActMax reference statistics
    over [start, stop) samples of the dataset.

    This step is expensive for large datasets (e.g., ~30 min for ImageNet + VGG16
    on a Titan RTX GPU). Results are cached to disk for fast retrieval later.

    Args:
        fv:         FeatureVisualization object.
        composite:  zennit LRP composite to use during backward pass.
        start:      Dataset start index.
        stop:       Dataset stop index (exclusive). None = len(dataset).
        batch_size: Number of samples per batch during precomputation.
    """
    if stop is None:
        stop = len(fv.dataset)

    print(f"[INFO] Starting precomputation for samples [{start}, {stop})...")
    print(f"[INFO] Batch size: {batch_size}")
    print(f"       This may take a while for large datasets.")

    fv.run(composite, start, stop, batch_size=batch_size)
    print("[INFO] Precomputation complete. Results cached to disk.")


def retrieve_max_references(
    fv: FeatureVisualization,
    concept_ids: list,
    layer_name: str,
    composite,
    mode: str = "relevance",
    top_n: int = 6,
) -> dict:
    """
    Retrieve reference images (RelMax or ActMax) for a set of concept IDs.

    Args:
        fv:          FeatureVisualization object (after fv.run() has been called).
        concept_ids: List of channel/concept indices to visualize.
        layer_name:  The layer name where these concepts reside.
        composite:   zennit composite (used to compute relevance maps for refs).
        mode:        'relevance' for RelMax, 'activation' for ActMax.
        top_n:       Number of top reference images per concept to retrieve.

    Returns:
        ref_c (dict): Dictionary mapping concept_id -> list of reference images/arrays.
    """
    print(f"[INFO] Retrieving '{mode}' references for concepts {concept_ids}")
    print(f"       Layer: '{layer_name}', top-{top_n} samples per concept.")

    ref_c = fv.get_max_reference(
        concept_ids=concept_ids,
        layer_name=layer_name,
        mode=mode,            # 'relevance' = RelMax, 'activation' = ActMax
        composite=composite,
        r_range=(0, top_n),   # indices of top references to retrieve
    )

    return ref_c


# ---------------------------------------------------------------------------
# Visualization
# ---------------------------------------------------------------------------

def visualize_references(ref_c: dict, title: str = "Concept References") -> None:
    """
    Display a grid of concept reference images using crp.image.plot_grid.

    Args:
        ref_c:  Reference image dict from fv.get_max_reference().
        title:  Optional title for the plot.
    """
    print(f"[INFO] Plotting reference grid: '{title}'")
    try:
        fig = plot_grid(ref_c)
        # Save to file instead of displaying (useful in headless environments)
        output_path = f"./{title.replace(' ', '_')}.png"
        fig.savef