diff --git a/rsseval/.github/raven-example.png b/rsseval/.github/raven-example.png new file mode 100644 index 00000000..9ae7c4e7 Binary files /dev/null and b/rsseval/.github/raven-example.png differ diff --git a/rsseval/README.md b/rsseval/README.md index 43c82726..98aa7272 100644 --- a/rsseval/README.md +++ b/rsseval/README.md @@ -67,6 +67,86 @@ In this setting, which is the same as the one presented in [Marconato et al. (20 ![Kandinsky pattern illustration](.github/kand-illustration.png) +## RAVEN + +This branch adds a reasoning-shortcut task based on the RAVEN dataset, +introduced by [Zhang et al. (2019)](https://arxiv.org/abs/1903.02741). RAVEN is +a computer-vision dataset inspired by John C. Raven's original 1938 Raven's +Progressive Matrices (RPM) test, a non-verbal assessment of abstract reasoning. +The dataset name is an homage to Raven; its purpose is to evaluate structural, +relational, and analogical visual reasoning. + +An RPM problem contains a 3 by 3 matrix of grayscale panels: the first eight +panels form the context, the bottom-right panel is missing, and the model must +select the candidate that completes the matrix. Each RAVEN problem therefore +contains 16 panels: eight context panels and eight answer candidates, with a +target index from `0` through `7`. + +![Example RAVEN matrix and answer candidates](.github/raven-example.png) + +*Example from [Zhang et al.'s RAVEN paper](https://openaccess.thecvf.com/content_CVPR_2019/html/Zhang_RAVEN_A_Dataset_for_Relational_and_Analogical_Visual_REasoNing_CVPR_2019_paper.html): the upper matrix is the context and the lower panels are the answer candidates.* + +The current implementation supports the reduced `center_single` +configuration. Every panel contains one centered object, and the model reasons +over three rule-governed attributes: **Type**, **Size**, and **Color**. +`ravendpl` combines a shared panel encoder with factorized DeepProbLog +reasoning, scoring the candidate that is consistent with the rules inferred +from the first two rows. + +### Dataset layout + +RAVEN data is not included in the repository. The reduced datasets are +generated with a modified RAVEN generator, available at +[jucamohedano/RAVEN](https://github.com/jucamohedano/RAVEN) (a fork of +[WellyZhang/RAVEN](https://github.com/WellyZhang/RAVEN)): + +- branch [`raven-3x3x3`](https://github.com/jucamohedano/RAVEN/tree/raven-3x3x3) + (commit `65c6ba9`) for the three-value dataset; +- branch [`raven-4x4x4`](https://github.com/jucamohedano/RAVEN/tree/raven-4x4x4) + (commit `8fd5eb8`) for the four-value dataset. + +Each branch README documents the reduced attribute domains. From the desired +branch, generate the dataset with the generator's Python 2.7 environment: + +```sh +python src/dataset/main.py --num-samples 5000 --save-dir /RAVEN-3x3x3 +``` + +The defaults (`--val 2 --test 2 --seed 1234`) produce a 6:2:2 split of +3000 train / 1000 val / 1000 test samples. Use `--save-dir +/RAVEN-4x4x4` on the `raven-4x4x4` branch for the four-value +dataset. Then place the generated `.npz` panel data and matching `.xml` +metadata files under `rss/data`. The default three-value dataset is expected at: + +```text +rss/data/RAVEN-3x3x3/center_single/ + RAVEN_*_train.npz + RAVEN_*_train.xml + RAVEN_*_val.npz + RAVEN_*_val.xml + RAVEN_*_test.npz + RAVEN_*_test.xml +``` + +The implementation also supports a four-value-per-attribute dataset in +`RAVEN-4x4x4`. Select it with `--n_values 4`; the documented command below +uses the default 3x3x3 setting. + +### Train RavenDPL + +From the `rss` directory, train the default RAVEN-3x3x3 task with: + +```sh +python main.py --dataset raven --model ravendpl --task raven \ + --raven_config center_single --n_epochs 100 --batch_size 32 \ + --lr 0.0001 --weight_decay 0.00001 --entropy --w_h 0.1 \ + --c_sup 0 --seed 123 --output_dir ../outputs/raven-3x3x3-seed-123 +``` + +`--c_sup` controls the fraction of training samples with concept labels, while +`--which_c` selects the modeled attributes to supervise: `0` for Type, `1` for +Size, and `2` for Color. Use `--which_c -1` to supervise all three attributes. + ## Structure of the code * The code structure is similar to [Marconato et al. (2024) bears](https://github.com/samuelebortolotti/bears): @@ -218,4 +298,4 @@ cd docs; make html ## Libraries and extra tools -This code is adapted from [Marconato et al. (2024) bears](https://github.com/samuelebortolotti/bears). \ No newline at end of file +This code is adapted from [Marconato et al. (2024) bears](https://github.com/samuelebortolotti/bears). diff --git a/rsseval/rss/backbones/raven_encoder.py b/rsseval/rss/backbones/raven_encoder.py new file mode 100644 index 00000000..84206ed8 --- /dev/null +++ b/rsseval/rss/backbones/raven_encoder.py @@ -0,0 +1,193 @@ +import torch +import torch.nn as nn + + +class RavenPanelEncoder(nn.Module): + """ + Encoder for a single RAVEN panel (160x160). + """ + def __init__(self, latent_dim=9): + super(RavenPanelEncoder, self).__init__() + self.latent_dim = latent_dim + self.backbone = nn.Sequential( + nn.Flatten(), + nn.Linear(160 * 160, 512), + nn.ReLU(), + nn.Linear(512, 128), + nn.ReLU(), + nn.Linear(128, self.latent_dim), + ) + + def forward(self, x): + return self.backbone(x) + + +class RavenMLP(nn.Module): + """ + Simple MLP for RAVEN 160x160 grayscale images. + Input: [batch, 16, 1, 160, 160] (processed one by one) + Output: [batch, 16, latent_dim] + + 3x3x3 attribute layout: + Type(3) | Size(3) | Color(3) = 9 total + """ + NAME = "RavenMLP" + + def __init__(self, latent_dim=9): + super(RavenMLP, self).__init__() + self.latent_dim = latent_dim + self.panel_encoder = RavenPanelEncoder(latent_dim=self.latent_dim) + + def forward(self, x): + # x shape: [batch, 16, 1, 160, 160] + + # Flatten batch and n_images for efficient parallel processing + b, n, c, h, w = x.shape + # Reshape to [B*N, C*H*W] compatible with Flatten in panel_encoder + # treat every single panel as an independent sample to pass it through the encoder + x_flat = x.reshape(b * n, -1) + + # Process all panels in one go (Efficient View) + z = self.panel_encoder(x_flat) # [b*16, latent_dim] + + # Reshape back to [batch, 16, latent_dim] + z = z.view(b, n, -1) + + # Return tuple (output, None) to match interface expected by DPL models + return z, None + +if __name__ == "__main__": + """ + Standalone training loop to verify that RavenMLP can learn to predict + concept attributes (Type, Size, Color) from 160x160 panel images. + + This is a *concept-only* sanity check — no answer-classification head. + We train with CrossEntropyLoss per attribute on the 8 context panels + and report per-attribute accuracy on the validation set. + """ + import sys, os + # Allow imports from the rss root when running as a script + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + + from torch.utils.data import DataLoader + from datasets.utils.raven_creation import RAVEN_Dataset + from utils.losses import RAVEN_Concept_Match + + # ── Hyperparameters ───────────────────────────────────────────── + BATCH_SIZE = 64 + LR = 1e-3 + EPOCHS = 20 + DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + N_VALUES = 3 # 3 for 3x3x3, 4 for 4x4x4 + LATENT_DIM = N_VALUES * 3 + DATA_PATH = f"data/RAVEN-{N_VALUES}x{N_VALUES}x{N_VALUES}" + + ATTR_NAMES = ["Type", "Size", "Color"] + ATTR_SLICES = [(0, N_VALUES), (N_VALUES, 2 * N_VALUES), (2 * N_VALUES, 3 * N_VALUES)] + + # ── Data ──────────────────────────────────────────────────────── + train_ds = RAVEN_Dataset(base_path=DATA_PATH, config="center_single", split="train") + val_ds = RAVEN_Dataset(base_path=DATA_PATH, config="center_single", split="val") + + train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True) + val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True) + + print(f"Train: {len(train_ds)} samples | Val: {len(val_ds)} samples") + print(f"Device: {DEVICE}") + + # ── Model & Optimizer ─────────────────────────────────────────── + model = RavenMLP(latent_dim=LATENT_DIM).to(DEVICE) + optimizer = torch.optim.Adam(model.parameters(), lr=LR) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS) + + n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"Model parameters: {n_params:,}") + + # ── Evaluation helper ─────────────────────────────────────────── + @torch.no_grad() + def evaluate(loader): + model.eval() + total_loss = 0.0 + correct = {name: 0 for name in ATTR_NAMES} + total = 0 + + for images, target, concepts in loader: + images = images.to(DEVICE) + concepts = concepts.to(DEVICE) + + z, _ = model(images) + + out_dict = {"CS": z, "CONCEPTS": concepts} + loss, _ = RAVEN_Concept_Match(out_dict) + total_loss += loss.item() * images.size(0) + + # Per-attribute accuracy on context panels (0-7) + z_ctx = z[:, :8] # [B, 8, 9] + c_ctx = concepts[:, :8] # [B, 8, 4] + + for i, (name, (lo, hi)) in enumerate(zip(ATTR_NAMES, ATTR_SLICES)): + preds = z_ctx[..., lo:hi].argmax(dim=-1) # [B, 8] + correct[name] += (preds == c_ctx[..., i]).sum().item() + + total += z_ctx.shape[0] * z_ctx.shape[1] # B * 8 + + avg_loss = total_loss / len(loader.dataset) + accs = {name: correct[name] / total * 100 for name in ATTR_NAMES} + return avg_loss, accs + + # ── Training loop ─────────────────────────────────────────────── + print(f"\n{'Epoch':>5} | {'Train Loss':>10} | {'Val Loss':>8} | " + + " | ".join(f"{n:>7}" for n in ATTR_NAMES) + + " | Avg Acc") + print("-" * 75) + + best_avg_acc = 0.0 + best_val_loss = float("inf") + patience = 3 + min_delta = 0.01 + wait = 0 + + for epoch in range(1, EPOCHS + 1): + model.train() + epoch_loss = 0.0 + + for images, target, concepts in train_loader: + images = images.to(DEVICE) + concepts = concepts.to(DEVICE) + + z, _ = model(images) + + out_dict = {"CS": z, "CONCEPTS": concepts} + loss, _ = RAVEN_Concept_Match(out_dict) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + epoch_loss += loss.item() * images.size(0) + + scheduler.step() + train_loss = epoch_loss / len(train_ds) + + # Validation + val_loss, val_accs = evaluate(val_loader) + avg_acc = sum(val_accs.values()) / len(val_accs) + + print(f"{epoch:5d} | {train_loss:10.4f} | {val_loss:8.4f} | " + + " | ".join(f"{val_accs[n]:6.2f}%" for n in ATTR_NAMES) + + f" | {avg_acc:5.2f}%") + + if avg_acc > best_avg_acc: + best_avg_acc = avg_acc + + # Early stopping on validation loss + if val_loss < best_val_loss - min_delta: + best_val_loss = val_loss + wait = 0 + else: + wait += 1 + if wait >= patience: + print(f"\nEarly stopping at epoch {epoch} (no val loss improvement for {patience} epochs)") + break + + print(f"\nBest average validation accuracy: {best_avg_acc:.2f}%") \ No newline at end of file diff --git a/rsseval/rss/datasets/raven.py b/rsseval/rss/datasets/raven.py new file mode 100644 index 00000000..ba4033fb --- /dev/null +++ b/rsseval/rss/datasets/raven.py @@ -0,0 +1,95 @@ + +from datasets.utils.base_dataset import BaseDataset +from datasets.utils.raven_creation import RAVEN_Dataset +from backbones.raven_encoder import RavenMLP +import time +from torch.utils.data import DataLoader + +# ── Concept labels per n_values ────────────────────────────────────── +_TYPE_LABELS = { + 3: ["Triangle", "Pentagon", "Circle"], + 4: ["Triangle", "Square", "Pentagon", "Circle"], +} +_SIZE_LABELS = { + 3: ["Small", "Medium", "Large"], + 4: ["Small", "Medium-small", "Medium-large", "Large"], +} +_COLOR_LABELS = { + 3: ["White", "Gray", "Black"], + 4: ["White", "Light-gray", "Dark-gray", "Black"], +} + +class RAVEN(BaseDataset): + NAME = "raven" + + def get_data_loaders(self): + start = time.time() + + config = getattr(self.args, 'raven_config', "center_single") + n = getattr(self.args, 'n_values', 3) + base_path = f"data/RAVEN-{n}x{n}x{n}" + + self.dataset_train = RAVEN_Dataset( + base_path=base_path, + config=config, + split="train", + c_sup=self.args.c_sup, + which_c=self.args.which_c, + ) + + self.dataset_val = RAVEN_Dataset( + base_path=base_path, + config=config, + split="val", + ) + + self.dataset_test = RAVEN_Dataset( + base_path=base_path, + config=config, + split="test", + ) + + print(f"Loaded datasets in {time.time()-start:.2f} s.") + self.print_stats() + + train_loader = DataLoader(self.dataset_train, batch_size=self.args.batch_size, shuffle=True, num_workers=4) + val_loader = DataLoader(self.dataset_val, batch_size=self.args.batch_size, shuffle=False, num_workers=4) + test_loader = DataLoader(self.dataset_test, batch_size=self.args.batch_size, shuffle=False, num_workers=4) + + return train_loader, val_loader, test_loader + + def get_backbone(self, args=None): + n = getattr(self.args, 'n_values', 3) if self.args else 3 + return RavenMLP(latent_dim=n * 3), None + + def get_split(self): + return 16, () + + def get_concept_labels(self): + n = getattr(self.args, 'n_values', 3) + return ["Type", "Size", "Color", "Number"], [ + _TYPE_LABELS.get(n, _TYPE_LABELS[3]), + _SIZE_LABELS.get(n, _SIZE_LABELS[3]), + _COLOR_LABELS.get(n, _COLOR_LABELS[3]), + [str(i) for i in range(9)], + ] + + def get_labels(self): + # 8 choices (0-7) + return [str(i) for i in range(8)] + + def print_stats(self): + print("## Statistics ##") + print("Train samples", len(self.dataset_train)) + print("Validation samples", len(self.dataset_val)) + print("Test samples", len(self.dataset_test)) + +if __name__ == "__main__": + from argparse import Namespace + dataset = RAVEN( + args=Namespace(batch_size=32) + ) + train_loader, val_loader, test_loader = dataset.get_data_loaders() + print(f"Training number of batches: {len(train_loader)}") + print(f"Validation number of batches: {len(val_loader)}") + print(f"Test number of batches: {len(test_loader)}") \ No newline at end of file diff --git a/rsseval/rss/datasets/utils/raven_creation.py b/rsseval/rss/datasets/utils/raven_creation.py new file mode 100644 index 00000000..38e1c64e --- /dev/null +++ b/rsseval/rss/datasets/utils/raven_creation.py @@ -0,0 +1,277 @@ + +import os +import glob +import numpy as np +import torch +import xml.etree.ElementTree as ET +from torch.utils.data import Dataset + +class RAVEN_Dataset(Dataset): + """ + RAVEN Dataset for RSBench. + Supports partial concept supervision. + """ + def __init__(self, base_path, config="center_single", split="train", + c_sup=1, which_c=[-1]): + self.base_path = base_path + self.config = config + self.split = split + self.c_sup = c_sup + self.which_c = which_c + self.is_train = split == "train" + + pattern = os.path.join(self.base_path, self.config, f"RAVEN_*_{self.split}.npz") + self.all_files = sorted(glob.glob(pattern)) + + if len(self.all_files) == 0: + print(f"Warning: No files found for {pattern}") + self.files = [] + else: + self.files = self.all_files + + # Deterministic supervision mask, analogous to other rsbench datasets. + rng = np.random.RandomState(0) + self.r_seq = rng.rand(len(self.files)) if len(self.files) > 0 else np.array([]) + + def __len__(self): + return len(self.files) + + def __getitem__(self, idx): + file_path = self.files[idx] + data = np.load(file_path) + + # Images: [16, 160, 160] + # Normalize to [0, 1] and add channel dimension: [16, 1, 160, 160] + images = data['image'].astype(np.float32) / 255.0 + images = torch.from_numpy(images).unsqueeze(1) + + # Target: 0-7 + target = torch.tensor(data['target'], dtype=torch.long) + + # Concepts: Extract attributes for all 16 panels (8 context + 8 choices) + # meta_matrix does NOT contain entity values (like Type=Triangle), only Rule activity. + # So we must parse XML to get ground truth concepts for supervision. + + concepts = self._extract_concepts_from_xml(file_path.replace('.npz', '.xml')) + + if self.is_train: + concepts = self._apply_concept_supervision_mask(concepts, idx) + + return images, target, concepts + + def _apply_concept_supervision_mask(self, concepts, idx): + """Mask concepts with -1 according to c_sup / which_c. + + We only supervise the first 3 modeled attributes (Type, Size, Color) + on the 8 context panels, mirroring RAVEN_Concept_Match. + """ + concepts = concepts.clone() + + # Sample-level supervision fraction. + if self.r_seq[idx] > self.c_sup: + concepts[:8, :3] = -1 + return concepts + + # Attribute-level supervision subset. + if not (len(self.which_c) == 1 and self.which_c[0] == -1): + for attr_idx in range(3): + if attr_idx not in self.which_c: + concepts[:8, attr_idx] = -1 + + return concepts + + def _extract_concepts_from_xml(self, xml_path): + """ + Extract concept values for all 16 panels. + Returns tensor of shape [16, 4] -> (Type, Size, Color, Number) + Values are indices 0-2 for Type/Size/Color (3x3x3 dataset). + """ + tree = ET.parse(xml_path) + root = tree.getroot() + + panels_data = [] + all_panels = root.findall('.//Panel') + + # 3x3x3 constrained dataset: + # Type: 1=triangle, 2=pentagon, 3=circle (in XML) -> 0-2 after -1 shift + # Size: 0=small(0.4), 1=medium(0.6), 2=large(0.9) + # Color: 0=white(255), 1=gray(140), 2=black(0) + + for panel in all_panels[:16]: + entity = panel.find('.//Entity') + layout = panel.find('.//Layout') + + # Defaults + p_c = [0, 0, 0, 0] # Type, Size, Color, Number + + if entity is not None: + try: p_c[0] = int(entity.get('Type', 0)) - 1 # Shift 1-3 to 0-2 + except: pass + try: p_c[1] = int(entity.get('Size', 0)) + except: pass + try: p_c[2] = int(entity.get('Color', 0)) + except: pass + + if layout is not None: + try: p_c[3] = int(layout.get('Number', 0)) + except: pass + + panels_data.append(p_c) + + return torch.tensor(panels_data, dtype=torch.long) + + + +if __name__ == "__main__": + import matplotlib.pyplot as plt + import matplotlib.patches as mpatches + from collections import Counter + + TYPE_LABELS = ["Triangle", "Square", "Pentagon", "Hexagon"] + SIZE_LABELS = ["Small", "Medium-S", "Medium-L", "Large"] + COLOR_LABELS = ["White", "Gray-W", "Gray-B", "Black"] + + dataset = RAVEN_Dataset( + base_path="data/RAVEN-4x4x4", + config="center_single", + split="train", + ) + print(f"Dataset size: {len(dataset)} samples") + + def extract_rules(xml_path): + """Parse elements from XML and return dict {attr: rule_name}.""" + tree = ET.parse(xml_path) + root = tree.getroot() + rules = {} + for rule_el in root.findall('.//Rule'): + attr = rule_el.get('attr', '') + name = rule_el.get('name', '') + rules[attr] = name + return rules + + # ── Fig 1: Three sample puzzles ────────────────────────────────── + n_samples = 3 + fig, axes = plt.subplots(n_samples, 5, figsize=(16, 3.5 * n_samples + 1), + gridspec_kw={"width_ratios": [3, 0.15, 2, 0.15, 2]}) + fig.suptitle("RAVEN-3×3×3 — Sample Puzzles", fontsize=15, fontweight="bold", y=0.98) + + for row in range(n_samples): + images, target, concepts = dataset[row] + imgs = images[:, 0].numpy() # [16, 160, 160] + target_idx = target.item() + c = concepts.numpy() # [16, 4] + + # Extract rules from XML + xml_path = dataset.files[row].replace('.npz', '.xml') + rules = extract_rules(xml_path) + + # --- Left: 3×3 context matrix (panel 9 = "?") --- + grid = np.ones((160 * 3 + 4, 160 * 3 + 4)) * 0.85 + for i in range(9): + r, col = divmod(i, 3) + y0 = r * (160 + 2) + x0 = col * (160 + 2) + if i < 8: + grid[y0:y0 + 160, x0:x0 + 160] = imgs[i] + else: + patch = np.ones((160, 160)) * 0.6 + # draw a "?" + grid[y0:y0 + 160, x0:x0 + 160] = patch + ax = axes[row, 0] + ax.imshow(grid, cmap="gray", vmin=0, vmax=1) + # Build rule annotation string + rule_parts = [] + for attr in ["Type", "Size", "Color"]: + rname = rules.get(attr, "?") + rule_parts.append(f"{attr}: {rname}") + rule_str = " | ".join(rule_parts) + ax.set_title(f"Sample {row}\n{rule_str}", fontsize=10, fontweight="bold") + ax.axis("off") + + # spacer + axes[row, 1].axis("off") + axes[row, 3].axis("off") + + # --- Middle: answer choices 0-3 --- + ans_grid = np.ones((160 * 2 + 2, 160 * 2 + 2)) * 0.85 + for j in range(4): + r, col = divmod(j, 2) + y0 = r * (160 + 2) + x0 = col * (160 + 2) + ans_grid[y0:y0 + 160, x0:x0 + 160] = imgs[8 + j] + ax = axes[row, 2] + ax.imshow(ans_grid, cmap="gray", vmin=0, vmax=1) + # highlight correct answer with a green rectangle + if target_idx < 4: + r, col = divmod(target_idx, 2) + rect = plt.Rectangle((col * 162 - 1, r * 162 - 1), 162, 162, + linewidth=3, edgecolor="limegreen", facecolor="none") + ax.add_patch(rect) + ax.set_title("Choices 0-3", fontsize=10) + ax.axis("off") + + # --- Right: answer choices 4-7 --- + ans_grid2 = np.ones((160 * 2 + 2, 160 * 2 + 2)) * 0.85 + for j in range(4): + r, col = divmod(j, 2) + y0 = r * (160 + 2) + x0 = col * (160 + 2) + ans_grid2[y0:y0 + 160, x0:x0 + 160] = imgs[12 + j] + ax = axes[row, 4] + ax.imshow(ans_grid2, cmap="gray", vmin=0, vmax=1) + if target_idx >= 4: + r, col = divmod(target_idx - 4, 2) + rect = plt.Rectangle((col * 162 - 1, r * 162 - 1), 162, 162, + linewidth=3, edgecolor="limegreen", facecolor="none") + ax.add_patch(rect) + ax.set_title("Choices 4-7", fontsize=10) + ax.axis("off") + + plt.tight_layout(rect=[0, 0, 1, 0.95]) + plt.savefig("raven_samples.png", dpi=150, bbox_inches="tight") + print("Saved raven_samples.png") + plt.show() + + # ── Fig 2: Attribute distributions ─────────────────────────────── + n_check = min(500, len(dataset)) + type_counts = Counter() + size_counts = Counter() + color_counts = Counter() + + for idx in range(n_check): + _, _, concepts = dataset[idx] + # Use context panels 0-7 + correct candidate (panel 8+target) + type_counts.update(concepts[:8, 0].tolist()) + size_counts.update(concepts[:8, 1].tolist()) + color_counts.update(concepts[:8, 2].tolist()) + + # ── Detect n_vals from the data ────────────────────────────── + n_type = max(type_counts.keys()) + 1 if type_counts else 3 + n_size = max(size_counts.keys()) + 1 if size_counts else 3 + n_color = max(color_counts.keys()) + 1 if color_counts else 3 + + fig, axes = plt.subplots(1, 3, figsize=(14, 4)) + fig.suptitle(f"Attribute Distributions (first {n_check} samples, context panels)", + fontsize=13, fontweight="bold") + + bar_colors = ["#4C72B0", "#55A868", "#C44E52", "#DD8452", "#937860"] + + for ax, counts, n_vals, name in [ + (axes[0], type_counts, n_type, "Type"), + (axes[1], size_counts, n_size, "Size"), + (axes[2], color_counts, n_color, "Color"), + ]: + labels = [str(i) for i in range(n_vals)] + vals = [counts.get(i, 0) for i in range(n_vals)] + colors = bar_colors[:n_vals] + bars = ax.bar(labels, vals, color=colors, edgecolor="black", linewidth=0.5) + ax.set_title(name, fontsize=12, fontweight="bold") + ax.set_ylabel("Count") + for bar, v in zip(bars, vals): + ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + max(vals) * 0.01, + str(v), ha="center", va="bottom", fontsize=9) + + plt.tight_layout() + plt.savefig("raven_distributions.png", dpi=150, bbox_inches="tight") + print("Saved raven_distributions.png") + plt.show() \ No newline at end of file diff --git a/rsseval/rss/models/__init__.py b/rsseval/rss/models/__init__.py index 0335c075..9cba9888 100644 --- a/rsseval/rss/models/__init__.py +++ b/rsseval/rss/models/__init__.py @@ -52,7 +52,8 @@ def get_model(args, encoder, decoder, n_images, c_split): "xordpl", "mnmathnn", "mnmathcbm", - "mnmathdpl" + "mnmathdpl", + "ravendpl" ]: return names[args.model]( encoder, n_images=n_images, c_split=c_split, args=args diff --git a/rsseval/rss/models/ravendpl.py b/rsseval/rss/models/ravendpl.py new file mode 100644 index 00000000..d920cb35 --- /dev/null +++ b/rsseval/rss/models/ravendpl.py @@ -0,0 +1,489 @@ +# RAVEN for DPL +import torch +import torch.nn as nn +import torch.nn.functional as F +from models.utils.deepproblog_modules import DeepProblogModel +from models.utils.utils_problog import build_worlds_queries_matrix_RAVEN +from models.utils.ops import outer_product +from utils.args import * +from utils.conf import get_device +from utils.dpl_loss import RAVEN_DPL +from utils.losses import RAVEN_Cumulative + + +def get_parser() -> ArgumentParser: + """Returns the parser + + Returns: + argparse: argument parser + """ + parser = ArgumentParser(description="Learning via" "Concept Extractor .") + add_management_args(parser) + add_experiment_args(parser) + return parser + + +class RavenDPL(DeepProblogModel): + """RAVEN DeepProbLog model with factorized logic per attribute.""" + + NAME = "ravendpl" + + def __init__( + self, + encoder, + n_images=16, + c_split=(), + model_dict=None, + n_facts=9, + nr_classes=8, + args=None, + ): + """Initialize method + + Args: + self: instance + encoder (nn.Module): encoder + n_images (int, default=16): number of images (8 context + 8 choices) + c_split: concept splits + model_dict (default=None): model dictionary + n_facts (int, default=9): number of concepts (3+3+3 for RAVEN-3x3x3) + nr_classes (int, default=8): number of choice candidates + args: command line arguments + + Returns: + None: This function does not return a value. + """ + super(RavenDPL, self).__init__( + encoder=encoder, + model_dict=model_dict, + n_facts=n_facts, + nr_classes=nr_classes, + ) + + # how many images and explicit split of concepts + self.n_images = n_images + self.c_split = c_split + + # Concept dimensions per attribute + self.raven_config = args.raven_config if args else "center_single" + if self.raven_config == "center_single": + n = getattr(args, 'n_values', 3) + self.dims = {"Type": n, "Size": n, "Color": n} + self.n_facts = n * 3 + + self.nr_classes = nr_classes # 8 candidate choices + self.n_and_classes = 2 # invalid / valid for explicit row-rule agreement + + # Per-attribute rules. + # In RAVEN-3x3x3 we also support Distribute_Three. For a 3-value domain, + # this means the row contains all three distinct values in some order. + # Arithmetic does not apply to Type. + self.attr_rules = {} + for attr in self.dims: + if attr == "Type": + self.attr_rules[attr] = ["Constant", "Progression", "Distribute_Three"] + else: + self.attr_rules[attr] = [ + "Constant", + "Progression", + "Arithmetic_Add", + "Arithmetic_Sub", + "Distribute_Three", + ] + + # Worlds-queries matrices (one per attribute, with attribute-specific rules) + # plus attribute-specific explicit row-agreement tables over rule pairs. + self.wq_matrices = {} + self.and_rules = {} + self.device = get_device() + # Generator uses offset=+1 for Size arithmetic, 0 for Color. + _ARITH_OFFSETS = {"Size": 1, "Color": 0, "Type": 0} + for attr, dim in self.dims.items(): + rules = self.attr_rules[attr] + offset = _ARITH_OFFSETS.get(attr, 0) + mat, and_rule = build_worlds_queries_matrix_RAVEN(dim, rules, arithmetic_offset=offset) + self.wq_matrices[attr] = mat.float().to(self.device) + self.and_rules[attr] = and_rule.float().to(self.device) + + + def forward(self, x): + """Forward method + + Args: + self: instance + x (torch.tensor): input vector [batch, 16, 1, 160, 160] + + Returns: + out_dict: output dictionary + """ + # 1. Encoding + z, _ = self.encoder(x) # [B, 16, n_facts] + + # 2. Concept Extraction (softmax + clamp per attribute) + concepts_probs = self.normalize_concepts(z) + + # Concatenate all attribute probs for supervision + extracted_probs = [concepts_probs[attr] for attr in self.dims] + pCS = torch.cat(extracted_probs, dim=-1) # [B, 16, n_facts] + + # 3. Factorized Symbolic Reasoning + # Compute row-level rule distributions for the two observed context rows. + row_rule_dists = {} + for attr, probs in concepts_probs.items(): + wq = self.wq_matrices[attr] + + # Row 1: panels [0, 1, 2] -> rule distribution [B, n_rules] + p_rule_r1 = self.problog_inference(probs[:, [0, 1, 2], :], wq) + # Row 2: panels [3, 4, 5] -> rule distribution [B, n_rules] + p_rule_r2 = self.problog_inference(probs[:, [3, 4, 5], :], wq) + + row_rule_dists[attr] = {"r1": p_rule_r1, "r2": p_rule_r2} + + # Score the 8 candidate choices via explicit 3-row agreement. + ys = self.score_choices(concepts_probs, row_rule_dists) + + # Prepare output dictionary + out_dict = {"YS": ys, "CS": z, "pCS": pCS} + + # Add row-level rule distributions for debugging/analysis. + for attr, row_dists in row_rule_dists.items(): + out_dict[f"{attr}_R1_PREDS"] = row_dists["r1"] + out_dict[f"{attr}_R2_PREDS"] = row_dists["r2"] + + return out_dict + + def normalize_concepts(self, z): + """Computes the probability for each ProbLog fact given the latent vector z, + applying softmax with epsilon clamping per attribute group. + + Args: + self: instance + z (torch.tensor): encoder output [B, 16, n_facts] + + Returns: + concepts_probs: dict mapping attribute name -> [B, 16, dim] probability tensor + """ + + def soft_clamp(h, dim=-1): + h = nn.Softmax(dim=dim)(h) + eps = 1e-5 + h = h + eps + with torch.no_grad(): + Z = torch.sum(h, dim=dim, keepdim=True) + h = h / Z + return h + + # Split z into per-attribute logits and apply softmax + clamp + concepts_probs = {} + offset = 0 + for attr, dim in self.dims.items(): + concepts_probs[attr] = soft_clamp(z[..., offset : offset + dim]) + offset += dim + + return concepts_probs + + def problog_inference(self, row_probs, wq): + """Problog inference for a single attribute over a single row of 3 panels. + + Computes the outer product of 3 panel concept distributions to build + the world probability vector, then maps to rule distributions via w_q. + + Args: + self: instance + row_probs (torch.tensor): [B, 3, n_vals] concept probs for 3 panels + wq (torch.tensor): [n_vals^3, n_rules] logic matrix + + Returns: + rule_dist: [B, n_rules] probability of each rule given this row + """ + # Number of rules is determined by the w_q matrix (varies per attribute) + n_rules = wq.shape[1] + + # Outer product of 3 panel distributions -> world probabilities + worlds_tensor = outer_product( + row_probs[:, 0], row_probs[:, 1], row_probs[:, 2] + ) # [B, N, N, N] + + worlds_prob = worlds_tensor.reshape(worlds_tensor.shape[0], -1) # [B, N^3] + + # Compute rule distribution: sum over worlds weighted by w_q + rule_dist = torch.zeros( + size=(worlds_prob.shape[0], n_rules), device=worlds_prob.device + ) + for i in range(n_rules): + rule_dist[:, i] = self.compute_query(i, worlds_prob, wq).view(-1) + + return rule_dist + + def problog_inference_log(self, row_log_probs, wq): + """Log-space ProbLog inference for a single attribute over a single row. + + Computes log P(rule) = logsumexp over worlds where the rule holds, + using log-space throughout to avoid gradient starvation from + near-uniform probability distributions. + + Args: + self: instance + row_log_probs (torch.tensor): [B, 3, n_vals] log concept probs for 3 panels + wq (torch.tensor): [n_vals^3, n_rules] binary logic matrix + + Returns: + log_rule_dist: [B, n_rules] log-probability of each rule given this row + """ + n_rules = wq.shape[1] + n_worlds = wq.shape[0] + batch_size = row_log_probs.shape[0] + + # Sum log-probs across 3 panels for each world (AND in log-space) + # row_log_probs: [B, 3, n_vals] + # worlds_log_prob[b, i*j*k] = row_log_probs[b,0,i] + row_log_probs[b,1,j] + row_log_probs[b,2,k] + log_p0 = row_log_probs[:, 0, :] # [B, n_vals] + log_p1 = row_log_probs[:, 1, :] # [B, n_vals] + log_p2 = row_log_probs[:, 2, :] # [B, n_vals] + + # Compute all world log-probs via broadcast + # [B, n_vals, 1, 1] + [B, 1, n_vals, 1] + [B, 1, 1, n_vals] -> [B, n_vals, n_vals, n_vals] + worlds_log = log_p0.unsqueeze(2).unsqueeze(3) + log_p1.unsqueeze(1).unsqueeze(3) + log_p2.unsqueeze(1).unsqueeze(2) + worlds_log = worlds_log.reshape(batch_size, n_worlds) # [B, n_vals^3] + + # Compute log rule distribution: logsumexp over worlds where rule holds + log_rule_dist = torch.full( + (batch_size, n_rules), -1e10, device=worlds_log.device + ) + for i in range(n_rules): + # Mask: which worlds have this rule + mask = wq[:, i].bool() # [n_worlds] + if mask.any(): + # logsumexp over the masked worlds + masked_log = worlds_log[:, mask] # [B, n_valid_worlds] + log_rule_dist[:, i] = torch.logsumexp(masked_log, dim=1) + + return log_rule_dist + + def compute_query(self, query, worlds_prob, wq): + """Computes query probability given the worlds probability P(w). + + Args: + self: instance + query (int): query index (rule index) + worlds_prob (torch.tensor): [B, N^3] world probabilities + wq (torch.tensor): [N^3, n_rules] logic matrix + + Returns: + query_prob: [B, 1] probability of the query + """ + # Select the column of w_q matrix corresponding to the current query + w_q = wq[:, query] + # Compute query probability by summing the probability of all worlds where the query is true + query_prob = torch.sum(w_q * worlds_prob, dim=1, keepdim=True) + return query_prob + + def score_choices(self, concepts_probs, row_rule_dists): + """Score the 8 candidate choices in log-space via explicit 3-row rule agreement. + + For each candidate and attribute, the model computes the log rule + distributions of row 1, row 2, and the candidate-completed row 3, + builds a joint world over ``(r1, r2, r3)``, and uses an attribute-specific + ``and_rule`` to compute the log-probability of agreement (r1==r2==r3). + + Log-space scoring avoids gradient starvation: when concept probabilities + are near-uniform, probability-space scores for all candidates are nearly + identical (differing by ~3e-5), producing vanishing gradients. In log + space, the same differences are amplified ~9x per attribute, keeping + the gradient signal alive through the reasoning layer. + + This mirrors how KandDPL (the rsbench Kand-Logic model) computes its + output: raw probabilities from the reasoning layer, with log applied in + the loss function. No intermediate softmax is applied. + + Args: + self: instance + concepts_probs: dict mapping attribute name -> [B, 16, dim] probs + row_rule_dists: dict mapping attribute name -> + {"r1": [B, n_rules], "r2": [B, n_rules]} (probability-space, kept for analysis) + + Returns: + ys: [B, 8] log-softmaxed choice log-probabilities + """ + # Compute log concept probabilities + concepts_log_probs = { + attr: probs.clamp(min=1e-8).log() + for attr, probs in concepts_probs.items() + } + + # Pre-compute log rule distributions for context rows (in log-space) + row_log_rule_dists = {} + for attr, log_probs in concepts_log_probs.items(): + wq = self.wq_matrices[attr] + log_rule_r1 = self.problog_inference_log( + log_probs[:, [0, 1, 2], :], wq + ) + log_rule_r2 = self.problog_inference_log( + log_probs[:, [3, 4, 5], :], wq + ) + row_log_rule_dists[attr] = {"r1": log_rule_r1, "r2": log_rule_r2} + + # Score each candidate + choice_log_scores = [] + + for c_idx in range(8, 16): + attr_log_scores = [] + + for attr, log_probs in concepts_log_probs.items(): + wq = self.wq_matrices[attr] + log_rule_r1 = row_log_rule_dists[attr]["r1"] + log_rule_r2 = row_log_rule_dists[attr]["r2"] + + # Candidate-completed Row 3 in log-space + log_rule_r3 = self.problog_inference_log( + log_probs[:, [6, 7, c_idx], :], wq + ) # [B, n_rules] + + # Compute log agreement probability + log_agreement = self.compute_three_row_rule_agreement_log( + log_rule_r1, log_rule_r2, log_rule_r3, attr + ) # [B] + + attr_log_scores.append(log_agreement) + + # Combine attribute log-scores by summing (AND across independent attributes) + total_log_score = torch.stack(attr_log_scores, dim=1).sum(dim=1) # [B] + choice_log_scores.append(total_log_score) + + log_scores = torch.stack(choice_log_scores, dim=1) # [B, 8] + + # Return log-softmax (no intermediate exp/softmax that kills gradients) + ys = F.log_softmax(log_scores, dim=-1) + + return ys + + def compute_three_row_rule_agreement_log(self, log_rule_r1, log_rule_r2, log_rule_r3, attr): + """Compute log-probability of agreement (r1==r2==r3) in log-space. + + For each valid rule triplet (r, r, r), computes: + log P(r1=r, r2=r, r3=r) = log_rule_r1[r] + log_rule_r2[r] + log_rule_r3[r] + Then marginalizes over all valid triplets via logsumexp. + + Args: + self: instance + log_rule_r1 (torch.tensor): [B, n_rules] log rule distribution for row 1 + log_rule_r2 (torch.tensor): [B, n_rules] log rule distribution for row 2 + log_rule_r3 (torch.tensor): [B, n_rules] log rule distribution for row 3 + attr (str): attribute name + + Returns: + log_agreement (torch.tensor): [B] log-probability of valid agreement + """ + and_rule = self.and_rules[attr] # [n_rules^3, 2] + # Valid triplets: where and_rule[:, 1] == 1 + valid_mask = and_rule[:, 1].bool() # [n_rules^3] + + if not valid_mask.any(): + return torch.full( + (log_rule_r1.shape[0],), -1e10, device=log_rule_r1.device + ) + + n_rules = log_rule_r1.shape[1] + + # Build log-probability for each rule triplet world + # log P(r1, r2, r3) = log_rule_r1[r1] + log_rule_r2[r2] + log_rule_r3[r3] + # Using broadcast: [B, n_rules, 1, 1] + [B, 1, n_rules, 1] + [B, 1, 1, n_rules] + log_triplet = ( + log_rule_r1.unsqueeze(2).unsqueeze(3) + + log_rule_r2.unsqueeze(1).unsqueeze(3) + + log_rule_r3.unsqueeze(1).unsqueeze(2) + ) # [B, n_rules, n_rules, n_rules] + log_triplet = log_triplet.reshape(log_rule_r1.shape[0], n_rules**3) # [B, n_rules^3] + + # Select only valid triplets and logsumexp + log_valid_triplets = log_triplet[:, valid_mask] # [B, n_valid] + log_agreement = torch.logsumexp(log_valid_triplets, dim=1) # [B] + + return log_agreement + + def compute_three_row_rule_agreement(self, p_rule_r1, p_rule_r2, p_rule_r3, attr): + """Compute explicit agreement over the latent rules of rows 1, 2, and 3. + + Kept for backward compatibility and analysis (confusion matrices, etc.). + The log-space version ``compute_three_row_rule_agreement_log`` is used + for training. + + Args: + self: instance + p_rule_r1 (torch.tensor): [B, n_rules] rule distribution for row 1 + p_rule_r2 (torch.tensor): [B, n_rules] rule distribution for row 2 + p_rule_r3 (torch.tensor): [B, n_rules] rule distribution for the + candidate-completed row 3 + attr (str): attribute name (`Type`, `Size`, or `Color`) + + Returns: + agreement_probs (torch.tensor): [B, 2] where column 0 is invalid + (`r1, r2, r3` not all equal) and column 1 is valid + (`r1 == r2 == r3`) + """ + n_rules = p_rule_r1.shape[1] + assert p_rule_r2.shape[1] == n_rules and p_rule_r3.shape[1] == n_rules, ( + p_rule_r1.shape, + p_rule_r2.shape, + p_rule_r3.shape, + ) + + rule_triplet_worlds = outer_product( + p_rule_r1, p_rule_r2, p_rule_r3 + ).reshape(-1, n_rules**3) + + agreement_probs = torch.zeros( + size=(p_rule_r1.shape[0], self.n_and_classes), device=p_rule_r1.device + ) + + and_rule = self.and_rules[attr] + for i in range(self.n_and_classes): + agreement_probs[:, i] = torch.sum( + and_rule[:, i] * rule_triplet_worlds, dim=1 + ) + + return agreement_probs + + + + @staticmethod + def get_loss(args): + """Returns the loss function + + Args: + args: command line arguments + + Returns: + loss: loss function + + Raises: + err: NotImplementedError if dataset is not specified + """ + if args.dataset == "raven": + return RAVEN_DPL(RAVEN_Cumulative) + else: + return NotImplementedError("Wrong dataset choice") + + def start_optim(self, args): + """Starts the optimizer + + Args: + self: instance + args: command line arguments + + Returns: + None: This function does not return a value. + """ + self.opt = torch.optim.Adam( + self.parameters(), + args.lr, + weight_decay=1e-5, + ) + + # override of to + def to(self, device): + super().to(device) + for attr in self.dims: + self.wq_matrices[attr] = self.wq_matrices[attr].to(device) + self.and_rules[attr] = self.and_rules[attr].to(device) + return self diff --git a/rsseval/rss/models/utils/utils_problog.py b/rsseval/rss/models/utils/utils_problog.py index b56164ca..ff163e8d 100644 --- a/rsseval/rss/models/utils/utils_problog.py +++ b/rsseval/rss/models/utils/utils_problog.py @@ -1022,4 +1022,88 @@ def create_mnist_and(sequence_len=0, n_digits=0, task="mnmath"): w_q[w, 1] = 1 return w_q else: - NotImplementedError("Wrong choice") \ No newline at end of file + NotImplementedError("Wrong choice") + + +def build_worlds_queries_matrix_RAVEN(n_vals, rules=["Constant"], arithmetic_offset=0): + """ + Build logic matrices for factorized RAVEN attribute experts. + + Args: + n_vals: number of possible values for the attribute. + rules: list of rule names to compile. + "Distribute_Three" is expanded into C(n_vals, 3) subset-specific + columns so that the AND-rule (r1==r2==r3) enforces that + all three rows use the *same* three values. + arithmetic_offset: bias added to the Arithmetic formula. + Generator uses 0 for Color, +1 for Size (v1+v2+1==v3). + + Returns: + w_q: tensor of shape [n_vals^3, n_rules] (with expanded columns) + and_rule: tensor of shape [n_rules^3, 2] encoding r1==r2==r3 + """ + from itertools import combinations, product as _product + + possible_worlds = list(_product(range(n_vals), repeat=3)) + n_worlds = len(possible_worlds) + + # Expand "Distribute_Three" into subset-specific columns. + expanded = [] + for rule in rules: + if rule == "Distribute_Three" and n_vals >= 3: + for s in combinations(range(n_vals), 3): + expanded.append(("Distribute_Three", s)) + elif rule == "Arithmetic_Add": + expanded.append(("Arithmetic_Add", arithmetic_offset)) + elif rule == "Arithmetic_Sub": + expanded.append(("Arithmetic_Sub", arithmetic_offset)) + else: + expanded.append((rule, None)) + n_rules = len(expanded) + + and_rule = torch.zeros(n_rules ** 3, 2) + for idx, (r1, r2, r3) in enumerate(_product(range(n_rules), repeat=3)): + if r1 == r2 and r2 == r3: + and_rule[idx, 1] = 1.0 + else: + and_rule[idx, 0] = 1.0 + + w_q = torch.zeros(n_worlds, n_rules) + + for w in range(n_worlds): + v1, v2, v3 = possible_worlds[w] + + for rule_idx, (rule, subset) in enumerate(expanded): + if rule == "Constant": + if v1 == v2 and v2 == v3: + w_q[w, rule_idx] = 1.0 + + elif rule == "Progression": + if (v2 - v1) == 1 and (v3 - v2) == 1: + w_q[w, rule_idx] = 1.0 + elif (v2 - v1) == -1 and (v3 - v2) == -1: + w_q[w, rule_idx] = 1.0 + elif (v2 - v1) == 2 and (v3 - v2) == 2: + w_q[w, rule_idx] = 1.0 + elif (v2 - v1) == -2 and (v3 - v2) == -2: + w_q[w, rule_idx] = 1.0 + + elif rule in ("Arithmetic", "Arithmetic_Add", "Arithmetic_Sub"): + # Split into separate columns so the AND-rule (r1==r2==r3) + # enforces the same operation across all three rows. + # Generator formulas differ per attribute: + # Color: v1+v2 == v3 | v1-v2 == v3 (offset=0) + # Size: v1+v2+1 == v3 | v1-v2-1 == v3 (offset=1) + o = arithmetic_offset + if rule in ("Arithmetic", "Arithmetic_Add"): + if v1 + v2 + o == v3: + w_q[w, rule_idx] = 1.0 + if rule in ("Arithmetic", "Arithmetic_Sub"): + if v1 - v2 - o == v3: + w_q[w, rule_idx] = 1.0 + + elif rule == "Distribute_Three": + if set([v1, v2, v3]) == set(subset): + w_q[w, rule_idx] = 1.0 + + return w_q, and_rule \ No newline at end of file diff --git a/rsseval/rss/utils/args.py b/rsseval/rss/utils/args.py index 23b464d7..abe04db0 100644 --- a/rsseval/rss/utils/args.py +++ b/rsseval/rss/utils/args.py @@ -37,10 +37,24 @@ def add_experiment_args(parser: ArgumentParser) -> None: "mini_patterns", "boia", "xor", - "mnmath" + "mnmath", + "raven", ], help="Which operation to choose.", ) + parser.add_argument( + "--raven_config", + type=str, + default="center_single", + help="RAVEN configuration (e.g. center_single)", + ) + parser.add_argument( + "--n_values", + type=int, + choices=[3, 4], + default=3, + help="Number of values per RAVEN attribute.", + ) # model settings parser.add_argument( "--model", @@ -128,6 +142,18 @@ def add_experiment_args(parser: ArgumentParser) -> None: "--n_epochs", type=int, default=50, help="Number of epochs per task." ) parser.add_argument("--batch_size", type=int, default=64, help="Batch size.") + parser.add_argument( + "--early_stop_patience", + type=int, + default=-1, + help="Stop if val loss does not improve for this many epochs (-1 disables).", + ) + parser.add_argument( + "--early_stop_min_delta", + type=float, + default=0.0, + help="Minimum val-loss decrease to count as an improvement.", + ) # deep ensembles parser.add_argument( @@ -204,7 +230,7 @@ def add_management_args(parser: ArgumentParser) -> None: "--wandb", type=str, default=None, - help="Enable wandb logging -- set name of project", + help="Enable wandb logging -- set the wandb entity/user (project is set via --project)", ) # checkpoints parser.add_argument( @@ -213,6 +239,12 @@ def add_management_args(parser: ArgumentParser) -> None: default=None, help="location and path FROM where to load ckpt.", ) + parser.add_argument( + "--output_dir", + type=str, + default=".", + help="Directory where training artifacts (best model, plots, CSVs) are saved.", + ) parser.add_argument( "--checkout", action="store_true", diff --git a/rsseval/rss/utils/checkpoint.py b/rsseval/rss/utils/checkpoint.py index 6b5b927c..8817a10b 100644 --- a/rsseval/rss/utils/checkpoint.py +++ b/rsseval/rss/utils/checkpoint.py @@ -4,6 +4,10 @@ from utils.conf import create_path +def _get_output_dir(args): + return getattr(args, "output_dir", ".") or "." + + def _get_tag(args): """Get tag for the model name @@ -31,12 +35,13 @@ def create_load_ckpt(model, args): Returns: model (nn.Module): model """ - create_path("data/runs") - create_path("data/ckpts") + output_dir = _get_output_dir(args) + create_path(os.path.join(output_dir, "data/runs")) + create_path(os.path.join(output_dir, "data/ckpts")) tag = _get_tag(args) - PATH = f"data/runs/{args.dataset}-{args.model}-{tag}-start.pt" + PATH = os.path.join(output_dir, f"data/runs/{args.dataset}-{args.model}-{tag}-start.pt") if args.checkin is not None: model.load_state_dict(torch.load(args.checkin)) @@ -60,10 +65,11 @@ def save_model(model, args): Returns: None: This function does not return a value. """ - create_path("data/ckpts") + output_dir = _get_output_dir(args) + create_path(os.path.join(output_dir, "data/ckpts")) tag = _get_tag(args) - PATH = f"data/ckpts/{args.dataset}-{args.model}-{tag}-{args.seed}-end.pt" + PATH = os.path.join(output_dir, f"data/ckpts/{args.dataset}-{args.model}-{tag}-{args.seed}-end.pt") if args.checkout: print("Saved", PATH, "\n") @@ -93,13 +99,14 @@ def load_checkpoint(model, args, checkin=None): Returns: model (nn.Module): model """ - create_path("data/ckpts") + output_dir = _get_output_dir(args) + create_path(os.path.join(output_dir, "data/ckpts")) tag = _get_tag(args) if checkin is not None: PATH = checkin else: - PATH = f"data/ckpts/{args.dataset}-{args.model}-{tag}-{args.seed}-end.pt" + PATH = os.path.join(output_dir, f"data/ckpts/{args.dataset}-{args.model}-{tag}-{args.seed}-end.pt") if not os.path.exists(PATH): raise ValueError(f"You have to train the model first, missing {PATH}") diff --git a/rsseval/rss/utils/dpl_loss.py b/rsseval/rss/utils/dpl_loss.py index 058bcf40..ab7ab453 100644 --- a/rsseval/rss/utils/dpl_loss.py +++ b/rsseval/rss/utils/dpl_loss.py @@ -172,6 +172,39 @@ def __init__(self, loss, nr_classes=2) -> None: self.base_loss = loss self.nr_classes = nr_classes + def forward(self, out_dict, args): + """Forward method + + Args: + self: instance + out_dict: output dictionary + args: command line arguments + + Returns: + loss: loss value + losses: losses dictionary + """ + loss, losses = self.base_loss(out_dict, args) + return loss, losses + + +class RAVEN_DPL(torch.nn.Module): + """RAVEN DPL loss""" + def __init__(self, loss, nr_classes=8) -> None: + """Initialize method + + Args: + self: instance + loss: loss function + nr_classes: number of classes + + Returns: + None: This function does not return a value. + """ + super().__init__() + self.base_loss = loss + self.nr_classes = nr_classes + def forward(self, out_dict, args): """Forward method diff --git a/rsseval/rss/utils/losses.py b/rsseval/rss/utils/losses.py index 9bd6e9ef..d93116e9 100644 --- a/rsseval/rss/utils/losses.py +++ b/rsseval/rss/utils/losses.py @@ -837,4 +837,146 @@ def MNMATH_Cumulative(out_dict: dict, args): mitigation += args.w_c * loss3 losses.update(losses3) + return loss + args.gamma * mitigation, losses + + +def RAVEN_Concept_Match(out_dict: dict): + """RAVEN concept supervision loss on the 8 context panels only. + + Applies per-attribute CrossEntropyLoss on the raw encoder logits (CS) + against ground-truth concept labels (CONCEPTS). Only the 8 context panels + (indices 0-7) are supervised; the 8 candidate panels are left unsupervised + so the model learns to score them through symbolic reasoning, not direct + concept supervision. + + Targets with value -1 are ignored, analogously to Kandinsky and the other + rsbench datasets supporting partial concept supervision. + + Attribute layout in z: Type(3) | Size(3) | Color(3) = 9 total + + Args: + out_dict: output dictionary containing: + "CS" [B, 16, n_facts]: raw encoder logits + "CONCEPTS" [B, 16, n_attrs]: ground-truth concept indices + + Returns: + loss: scalar loss value (averaged over supervised attributes) + losses: dictionary with "c-loss" entry + """ + z = out_dict["CS"][:, :8] # [B, 8, n_facts] + targets = out_dict["CONCEPTS"][:, :8].to(torch.long) # [B, 8, n_attrs] + + n = z.shape[-1] // 3 # values per attribute + loss = torch.tensor(0.0, device=z.device) + n_supervised_attrs = 0 + + for attr_idx, (lo, hi) in enumerate([(0, n), (n, 2 * n), (2 * n, 3 * n)]): + target = targets[..., attr_idx].reshape(-1) + mask = target != -1 + if mask.sum() > 0: + logits = z[..., lo:hi].reshape(-1, n) + loss += torch.nn.CrossEntropyLoss()(logits[mask], target[mask]) + n_supervised_attrs += 1 + + if n_supervised_attrs > 0: + loss /= n_supervised_attrs + + return loss, {"c-loss": loss.item()} + + +def RAVEN_Entropy(out_dict, args): + """RAVEN entropy loss. + + Follows the MiniKandinsky/Kandinsky structure, but applies it to the + 8 context panels only. Each context panel contributes 3 concept slots + (Type, Size, Color), each of size 3. + + Args: + out_dict: output dictionary containing "pCS" [B, 16, 9] + args: command line arguments + + Returns: + loss: entropy penalty value + losses: dictionary with "H-loss" entry + """ + # Only regularize the context panels, mirroring concept supervision. + pCs = out_dict["pCS"][:, :8] # [B, 8, n_facts] + n = pCs.shape[-1] // 3 # values per attribute + + # Split [Type(n), Size(n), Color(n)] and flatten panel/attribute slots. + pc_i = torch.cat(torch.split(pCs, n, dim=-1), dim=1) # [B, 24, n] + + # Mean predicted concept distribution across the batch for each slot. + p_mean = torch.mean(pc_i, dim=0) # [24, 3] + + p_mean += 1e-5 + # renormalization per slot distribution + with torch.no_grad(): + Z = torch.sum(p_mean, dim=1, keepdim=True) + p_mean /= Z + + loss = 0 + for i in range(p_mean.size(0)): + # loss -= torch.sum(p_mean[i] * p_mean[i].log()) / np.log(10) / p_mean.size(0) + loss -= torch.sum(p_mean[i] * p_mean[i].log()) / np.log(p_mean.size(1)) / p_mean.size(0) + + losses = {"H-loss": 1 - loss} + + assert (1 - loss) > 0, loss + + return 1 - loss, losses + + +def RAVEN_Classification(out_dict: dict, args): + """RAVEN classification loss. + + YS is a log-softmaxed [B, 8] log-probability distribution over the 8 + candidates (returned by score_choices in log-space). We apply NLL loss + directly on the log-probabilities, matching how KandDPL handles its output. + + Args: + out_dict: output dictionary containing "YS" [B, 8] log-probs and "LABELS" [B] + args: command line arguments + + Returns: + loss: scalar loss value + losses: dictionary with "y-loss" entry + """ + out = out_dict["YS"] # [B, 8] log-softmax choice log-probabilities + labels = out_dict["LABELS"].to(torch.long) # [B] ground-truth choice index (0-7) + + # YS is already log-probabilities, so NLL loss applies directly. + # No need for .log() (which would be log-log) or .clamp(). + loss = F.nll_loss(out, labels, reduction="mean") + + assert loss > 0, f"{loss}, {out}, {labels}" + + losses = {"y-loss": loss.item()} + return loss, losses + + +def RAVEN_Cumulative(out_dict: dict, args): + """RAVEN cumulative loss: task classification + optional entropy + optional concept supervision.""" + loss, losses = RAVEN_Classification(out_dict, args) + + mitigation = 0 + if args.entropy: + # Entropy annealing: if entropy_anneal_epochs > 0, linearly decay w_h + # from its full value down to 0.1 over the specified number of epochs, + # then hold at 0.1. This gives the task signal room to compete once + # concepts are bootstrapped. + w_h = args.w_h + anneal = getattr(args, 'entropy_anneal_epochs', 0) + if anneal > 0: + epoch = getattr(args, '_current_epoch', 0) + frac = max(0.0, 1.0 - epoch / anneal) + w_h = args.w_h * frac + 0.1 * (1.0 - frac) + loss_h, losses_h = RAVEN_Entropy(out_dict, args) + mitigation += w_h * loss_h + losses.update(losses_h) + if args.c_sup > 0: + loss_c, losses_c = RAVEN_Concept_Match(out_dict) + mitigation += args.w_c * loss_c + losses.update(losses_c) + return loss + args.gamma * mitigation, losses \ No newline at end of file diff --git a/rsseval/rss/utils/metrics.py b/rsseval/rss/utils/metrics.py index 80bd13a2..eaad2b23 100644 --- a/rsseval/rss/utils/metrics.py +++ b/rsseval/rss/utils/metrics.py @@ -8,7 +8,7 @@ import torch.nn as nn import torch import torch.nn.functional as F -from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score +from sklearn.metrics import accuracy_score, confusion_matrix, f1_score from scipy.special import softmax @@ -193,7 +193,11 @@ def evaluate_metrics( "mnmath", ]: loss, ac, acc, f1 = MNMATH_eval_tloss_cacc_acc(out_dict, concepts) + elif args.dataset == "raven": + loss, ac, acc, f1, rcf1 = RAVEN_eval_tloss_cacc_acc(out_dict, concepts, cf1=True) + fcf1 += rcf1 else: + print(f"Metrics not implemented for dataset: {args.dataset}") NotImplementedError() if not last: @@ -228,6 +232,27 @@ def evaluate_metrics( cs = np.split(c_pred, c_pred.shape[1], axis=1) p_cs = np.split(pc_pred, pc_pred.shape[1], axis=1) p_ys = y_pred + elif args.dataset == "raven": + # YS: [Batch, 8] -> [Batch] (candidate index) + y_true = y_true.astype(int).reshape(-1) + ys = np.argmax(y_pred, axis=1).astype(int).reshape(-1) + + p_cs_all = pc_pred + + # Derive per-attribute dimension from the tensor shape itself. + n = pc_pred.shape[-1] // 3 # values per attribute (3 or 4) + ptype = pc_pred[..., 0:n].argmax(axis=-1) + psize = pc_pred[..., n:2*n].argmax(axis=-1) + pcolor = pc_pred[..., 2*n:3*n].argmax(axis=-1) + + ctype_max = pc_pred[..., 0:n].max(axis=-1) + csize_max = pc_pred[..., n:2*n].max(axis=-1) + ccolor_max = pc_pred[..., 2*n:3*n].max(axis=-1) + + cs = np.stack([ptype, psize, pcolor], axis=-1) + p_cs = np.stack([ctype_max, csize_max, ccolor_max], axis=-1) + gs = c_true[..., :3] + p_ys = y_pred.max(axis=1) else: ys = np.argmax(y_pred, axis=1) @@ -236,7 +261,12 @@ def evaluate_metrics( p_cs = np.split(pc_pred, pc_pred.shape[1], axis=1) p_ys = y_pred - p_cs_all = p_cs + # For RAVEN, p_cs_all must remain the full concept-probability tensor + # [N, 16, 9] = [Type(3) | Size(3) | Color(3)] per panel. + # Do not overwrite it with p_cs, which only stores max confidences + # [N, 16, 3]. Other datasets keep the historical behavior. + if args.dataset != "raven": + p_cs_all = p_cs p_ys_all = y_pred assert len(gs) == len(cs), f"gs: {gs.shape}, cs: {cs.shape}" @@ -247,6 +277,7 @@ def evaluate_metrics( "presddoia", "clipboia", "clipsddoia", + "raven", ]: gs = np.concatenate(gs, axis=0).squeeze(1) if args.dataset in [ @@ -257,6 +288,8 @@ def evaluate_metrics( "clipsddoia", ]: cs = (cs >= 0.5).astype(np.int) + elif args.dataset == "raven": + pass elif args.dataset not in [ "kandinsky", "prekandinsky", @@ -288,6 +321,7 @@ def evaluate_metrics( "presddoia", "clipboia", "clipsddoia", + "raven", ]: p_cs_all = np.concatenate(p_cs_all, axis=0).squeeze( 1 @@ -321,6 +355,7 @@ def evaluate_metrics( "restrictedmnist", "clipsddoia", "clipshortmnist", + "raven", ]: if cf1: return tloss / L, cacc / L, yacc / L, f1sc / L, fcf1 / L @@ -1299,3 +1334,237 @@ def world_accuracy(world_prob: ndarray, world_true: ndarray, n_concepts: int): ).astype(int) return get_accuracy_and_counter(n_world, world_pred, world_true, True) + + +def raven_joint_concept_collapse(c_true, c_pred, n_vals=None): + """Joint concept collapse over the (Type, Size, Color) product space. + + Encodes (T,S,C) as code = T*n_vals^2 + S*n_vals + C, computes a + full-grid (n_vals^3 x n_vals^3) confusion matrix, and returns + 1 - coverage (higher = more collapse) and the matrix itself. + + Args: + c_true: (N, 3) integer array, columns = [Type, Size, Color] + c_pred: (N, 3) integer array, same layout + n_vals: cardinality per attribute. Auto-detected from data if None. + + Returns: + collapse: float in [0, 1] + cm: (n_vals^3, n_vals^3) confusion matrix + """ + c_true = np.asarray(c_true) + c_pred = np.asarray(c_pred) + + if n_vals is None: + n_vals = int(max(c_true.max(), c_pred.max())) + 1 + + assert c_true.shape[1] == 3, f"Expected (N, 3) with columns [Type, Size, Color]; got {c_true.shape}" + + # Drop rows where any ground-truth attribute is masked (-1). + # Without this, -1 values produce negative codes that confusion_matrix + # silently drops, losing data without warning. + valid = (c_true >= 0).all(axis=1) & (c_true < n_vals).all(axis=1) + c_true = c_true[valid] + c_pred = c_pred[valid] + + code_t = c_true[:, 0] * n_vals ** 2 + c_true[:, 1] * n_vals + c_true[:, 2] + code_p = c_pred[:, 0] * n_vals ** 2 + c_pred[:, 1] * n_vals + c_pred[:, 2] + + labels = list(range(n_vals ** 3)) + cm = confusion_matrix(code_t, code_p, labels=labels) + + max_per_col = np.max(cm, axis=0) + coverage = np.sum(np.clip(max_per_col, 0, 1)) / len(max_per_col) + return 1.0 - coverage, cm + + +def raven_pairwise_joint_collapse(c_true, c_pred, n_vals=None): + """Pairwise joint concept collapse for all attribute pairs. + + Computes three (n_vals^2) x (n_vals^2) confusion matrices (TypexSize, + TypexColor, SizexColor) and returns collapse = 1 - coverage for each. + + Args: + c_true: (N, 3) integer array, columns = [Type, Size, Color] + c_pred: (N, 3) integer array, same layout + n_vals: cardinality per attribute. Auto-detected from data if None. + + Returns: + dict mapping pair name to (collapse, cm) tuples + """ + c_true = np.asarray(c_true) + c_pred = np.asarray(c_pred) + assert c_true.shape[1] == 3, f"Expected (N, 3); got {c_true.shape}" + + if n_vals is None: + n_vals = int(max(c_true.max(), c_pred.max())) + 1 + + valid = (c_true >= 0).all(axis=1) & (c_true < n_vals).all(axis=1) + c_true = c_true[valid] + c_pred = c_pred[valid] + + pairs = [(0, 1, "TypexSize"), (0, 2, "TypexColor"), (1, 2, "SizexColor")] + n2 = n_vals ** 2 + labels = list(range(n2)) + result = {} + for i, j, name in pairs: + code_t = c_true[:, i] * n_vals + c_true[:, j] + code_p = c_pred[:, i] * n_vals + c_pred[:, j] + cm = confusion_matrix(code_t, code_p, labels=labels) + max_per_col = np.max(cm, axis=0) + coverage = np.sum(np.clip(max_per_col, 0, 1)) / len(max_per_col) + result[name] = (1.0 - coverage, cm) + return result + + +def plot_raven_pairwise_confusion_matrix(cm, attr_i, attr_j, n_vals=None, save_path=None): + """Plot a pairwise joint confusion matrix.""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + if n_vals is None: + n_vals = int(round(cm.shape[0] ** 0.5)) + n = n_vals ** 2 + assert cm.shape == (n, n), f"Expected ({n},{n}), got {cm.shape}" + + attr_names = ["T", "S", "C"] + tick_labels = [f"{attr_names[attr_i]}{c // n_vals}{attr_names[attr_j]}{c % n_vals}" for c in range(n)] + + fig, ax = plt.subplots(figsize=(8, 7)) + im = ax.imshow(cm + 1e-6, interpolation="nearest", cmap=plt.cm.Blues, + norm=matplotlib.colors.LogNorm()) + ax.set_title(f"Joint confusion: {attr_names[attr_i]}x{attr_names[attr_j]}") + fig.colorbar(im, ax=ax) + + ax.set_xticks(np.arange(n)) + ax.set_yticks(np.arange(n)) + ax.set_xticklabels(tick_labels, rotation=90, fontsize=9) + ax.set_yticklabels(tick_labels, fontsize=9) + ax.set_xlabel(f"Predicted ({attr_names[attr_i]},{attr_names[attr_j]})") + ax.set_ylabel(f"True ({attr_names[attr_i]},{attr_names[attr_j]})") + + ax.set_xticks(np.arange(n + 1) - 0.5, minor=True) + ax.set_yticks(np.arange(n + 1) - 0.5, minor=True) + ax.grid(which="minor", color="lightgray", linestyle="-", linewidth=0.3) + ax.tick_params(which="minor", bottom=False, left=False) + + # Annotate cells with counts + for i in range(n): + for j in range(n): + val = cm[i, j] + if val > 0: + ax.text(j, i, str(int(val)), ha="center", va="center", + fontsize=7, color="white" if val > cm.max() / 2 else "black") + + plt.tight_layout() + if save_path: + plt.savefig(save_path, dpi=200, bbox_inches="tight") + plt.close() + + +def plot_raven_joint_confusion_matrix(cm, n_vals=None, save_path=None, title="Joint concept confusion (T,S,C)"): + """Plot the full-grid joint (Type, Size, Color) confusion matrix.""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + if n_vals is None: + n_vals = int(round(cm.shape[0] ** (1/3))) + n = n_vals ** 3 + assert cm.shape == (n, n), f"Expected ({n},{n}), got {cm.shape}" + + tick_labels = [ + f"T{c // n_vals**2}S{(c // n_vals) % n_vals}C{c % n_vals}" + for c in range(n) + ] + + fig, ax = plt.subplots(figsize=(14, 12)) + im = ax.imshow(cm + 1e-6, interpolation="nearest", cmap=plt.cm.Blues, + norm=matplotlib.colors.LogNorm()) + ax.set_title(title) + fig.colorbar(im, ax=ax) + + ax.set_xticks(np.arange(n)) + ax.set_yticks(np.arange(n)) + ax.set_xticklabels(tick_labels, rotation=90, fontsize=7) + ax.set_yticklabels(tick_labels, fontsize=7) + ax.set_xlabel("Predicted (T,S,C)") + ax.set_ylabel("True (T,S,C)") + + ax.set_xticks(np.arange(n + 1) - 0.5, minor=True) + ax.set_yticks(np.arange(n + 1) - 0.5, minor=True) + ax.grid(which="minor", color="lightgray", linestyle="-", linewidth=0.3) + ax.tick_params(which="minor", bottom=False, left=False) + + plt.tight_layout() + if save_path: + plt.savefig(save_path, dpi=200, bbox_inches="tight") + plt.close() + + +def RAVEN_eval_tloss_cacc_acc(out_dict, concepts, cf1=False): + """RAVEN evaluation for the current 3x3x3 factorized setup. + + Per-attribute concept accuracy and F1, plus label accuracy and F1. + + Args: + out_dict: dictionary of model outputs + concepts: ground-truth concept tensor + cf1: if True, also return concept macro F1 as 5th value + + Returns: + loss: NLL loss on answer prediction + cacc: mean concept accuracy across attributes (%) + acc: label accuracy (%) + f1: label macro F1 (%) + cf1: mean concept macro F1 across attributes (%), only if cf1=True + """ + pCs = out_dict["pCS"] # [B, 16, n_facts] + n = pCs.shape[-1] // 3 # values per attribute + + # Type(0:n), Size(n:2n), Color(2n:3n) + ptype = pCs[..., 0:n].argmax(dim=-1) + psize = pCs[..., n:2*n].argmax(dim=-1) + pcolor = pCs[..., 2*n:3*n].argmax(dim=-1) + + g_type = concepts[..., 0] + g_size = concepts[..., 1] + g_color = concepts[..., 2] + + # Flatten panels before metric computation. If concept supervision masks + # are present, ignore masked slots (-1). + pred_attrs = [ptype.reshape(-1), psize.reshape(-1), pcolor.reshape(-1)] + true_attrs = [g_type.reshape(-1), g_size.reshape(-1), g_color.reshape(-1)] + + accs, f1s = [], [] + for y_true, y_pred in zip(true_attrs, pred_attrs): + mask = y_true != -1 + y_true = y_true[mask] + y_pred = y_pred[mask] + accs.append((y_pred == y_true).float().mean().item()) + f1s.append( + f1_score( + y_true.detach().cpu().numpy(), + y_pred.detach().cpu().numpy(), + average="macro", + ) + ) + + cacc = float(np.mean(accs)) * 100.0 + rcf1 = float(np.mean(f1s)) * 100.0 + + ys = out_dict["YS"] + labels = out_dict["LABELS"] + preds = ys.argmax(dim=-1) + acc = (preds == labels).float().mean().item() * 100.0 + f1 = f1_score(labels.cpu().numpy(), preds.cpu().numpy(), average="macro") * 100.0 + + # Real evaluation loss for the answer prediction. + # ys is already log-probabilities (log_softmax), so NLL applies directly. + loss = F.nll_loss(ys, labels.to(torch.long), reduction="mean") + + if cf1: + return loss, cacc, acc, f1, rcf1 + else: + return loss, cacc, acc, f1 diff --git a/rsseval/rss/utils/train.py b/rsseval/rss/utils/train.py index 635bafb7..8df9bb77 100644 --- a/rsseval/rss/utils/train.py +++ b/rsseval/rss/utils/train.py @@ -18,14 +18,20 @@ evaluate_mix, mean_entropy, accuracy_binary, + raven_joint_concept_collapse, + plot_raven_joint_confusion_matrix, + raven_pairwise_joint_collapse, + plot_raven_pairwise_confusion_matrix, ) from utils.generative import conditional_gen, recon_visaulization from utils import fprint import matplotlib.pyplot as plt from warmup_scheduler import GradualWarmupScheduler -from sklearn.metrics import multilabel_confusion_matrix, confusion_matrix -import numpy as np +from sklearn.metrics import ( + multilabel_confusion_matrix, + confusion_matrix, +) def convert_to_categories(elements): @@ -285,6 +291,13 @@ def save_predictions_to_csv(model, test_set, csv_name, dataset): elif "mnmath" in dataset: cs = torch.argmax(cs, dim=2) cs_true = cs_true.reshape(cs_true.size(0), cs_true.size(1) * cs_true.size(2)) + elif "raven" in dataset: + if y_true.dim() == 1: + y_true = y_true.unsqueeze(1) + if cs.dim() > 2: + cs = cs.reshape(cs.shape[0], -1) + if cs_true.dim() > 2: + cs_true = cs_true.reshape(cs_true.shape[0], -1) concatenated_tensor = ( torch.concatenate((ys, y_true, cs, cs_true), dim=1).cpu().detach().numpy() @@ -311,11 +324,21 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): None: This function does not return a value. """ + output_dir = getattr(args, "output_dir", ".") or "." + os.makedirs(output_dir, exist_ok=True) + + def out_path(name): + return os.path.join(output_dir, name) + # name - csv_name = f"{args.dataset}-{args.model}-lr-{args.lr}.csv" + csv_name = out_path(f"{args.dataset}-{args.model}-lr-{args.lr}.csv") # best f1 best_f1 = 0.0 + best_tloss = float("inf") + early_stop_wait = 0 + early_stop_patience = int(getattr(args, "early_stop_patience", -1)) + early_stop_min_delta = float(getattr(args, "early_stop_min_delta", 0.0)) to_add = "" if args.model in ["kandcbm", "sddoiacbm", "boiacbm", "mnistcbm"]: @@ -324,7 +347,7 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): if args.dataset in ["shortmnist"] and args.joint: to_add += "_joint" - save_path = f"best_model_{args.dataset}_{args.model}_{args.seed}{to_add}.pth" + save_path = out_path(f"best_model_{args.dataset}_{args.model}_{args.seed}{to_add}.pth") # save embeddings variable save_embeddings_flag = False @@ -352,7 +375,7 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): wandb.init( project=args.project, entity=args.wandb, - name=str(args.dataset) + "_" + str(args.model), + name=getattr(args, 'run_name', f"{args.dataset}_{args.model}"), config=args, ) @@ -368,6 +391,7 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): conc_sup = dataset.get_sup() for epoch in range(args.n_epochs): + args._current_epoch = epoch model.train() ys, y_true, cs, cs_true = None, None, None, None @@ -500,6 +524,22 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): lr=float(scheduler.get_last_lr()[0]), ) + # Early stopping based on validation loss. + if early_stop_patience > 0: + current_tloss = float(tloss) + if (best_tloss - current_tloss) > early_stop_min_delta: + best_tloss = current_tloss + early_stop_wait = 0 + else: + early_stop_wait += 1 + if early_stop_wait >= early_stop_patience: + print( + f"Early stopping at epoch {epoch}: " + f"val loss did not improve for {early_stop_patience} epochs " + f"(best={best_tloss:.6f}, current={current_tloss:.6f})." + ) + break + if args.dataset in ["clipshortmnist", "shortmnist"]: pass elif not args.tuning: @@ -550,13 +590,13 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): ] plot_multilabel_confusion_matrix( - y_true, y_pred, y_labels, "Labels", save_path="labels.png" + y_true, y_pred, y_labels, "Labels", save_path=out_path("labels.png") ) cfs = plot_actions_confusion_matrix( - c_true, c_pred, "Concepts", save_path="total_concepts_" + c_true, c_pred, "Concepts", save_path=out_path("total_concepts_") ) cf = plot_multilabel_confusion_matrix( - c_true, c_pred, concept_labels, "Concepts", save_path="total_concepts" + c_true, c_pred, concept_labels, "Concepts", save_path=out_path("total_concepts") ) print("Concept collapse", 1 - compute_coverage(cf)) @@ -570,17 +610,65 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): ["{i}" for i in range(10) for _ in range(4)] ] plot_multilabel_confusion_matrix( - y_true, y_pred, y_labels, "Labels", save_path="labels.png" + y_true, y_pred, y_labels, "Labels", save_path=out_path("labels.png") ) cf = plot_confusion_matrix( c_true, c_pred, labels=dataset.get_concept_labels(), title="Concepts", - save_path=f"concepts_{args.dataset}_{args.model}_lr_{args.lr}.png", + save_path=out_path(f"concepts_{args.dataset}_{args.model}_lr_{args.lr}.png"), ) print("Concept collapse", 1 - compute_coverage(cf)) + elif args.task == "raven": + plot_confusion_matrix( + y_true, + y_pred, + labels=dataset.get_labels(), + title="Labels", + save_path=out_path(f"labels_{args.dataset}_{args.model}_lr_{args.lr}.png"), + ) + + concept_names, concept_values = dataset.get_concept_labels() + + # Flatten [Batch, 16, Attr] -> [Total_Panels, Attr] + c_true_flat = c_true.reshape(-1, c_true.shape[-1]) + c_pred_flat = c_pred.reshape(-1, c_pred.shape[-1]) + + for i in range(c_pred_flat.shape[1]): + attr_name = concept_names[i] + attr_labels = concept_values[i] + plot_confusion_matrix( + c_true_flat[:, i], + c_pred_flat[:, i], + labels=attr_labels, + title=f"Concepts - {attr_name}", + save_path=out_path(f"concepts_{attr_name}_{args.dataset}_{args.model}_lr_{args.lr}.png"), + ) + cf_attr=confusion_matrix(c_true_flat[:, i], c_pred_flat[:, i], labels=list(range(len(attr_labels)))) + print(f"Concept collapse {attr_name}", 1 - compute_coverage(cf_attr)) + + # Joint (Type, Size, Color) concept collapse + joint_collapse, joint_cm = raven_joint_concept_collapse(c_true_flat, c_pred_flat) + n = int(round(joint_cm.shape[0] ** (1/3))) # derive n_vals from matrix size + print(f"Joint concept collapse ({n**3}-class): {joint_collapse:.4f}") + plot_raven_joint_confusion_matrix( + joint_cm, + save_path=out_path(f"joint_concepts_{args.dataset}_{args.model}_lr_{args.lr}.png"), + ) + + # Pairwise joint collapse + pairwise = raven_pairwise_joint_collapse(c_true_flat, c_pred_flat) + pair_idx = {"TypexSize": (0, 1), "TypexColor": (0, 2), "SizexColor": (1, 2)} + for pair_name, (clp, cm) in pairwise.items(): + i, j = pair_idx[pair_name] + safe_name = pair_name.replace("×", "x") + print(f"Pairwise collapse {pair_name} (9-class): {clp:.4f}") + plot_raven_pairwise_confusion_matrix( + cm, i, j, + save_path=out_path(f"pairwise_{safe_name}_{args.dataset}_{args.model}_lr_{args.lr}.png"), + ) else: if args.task in ["patterns", "mini_patterns"]: @@ -592,7 +680,7 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): y_pred, labels=dataset.get_labels(), title="Labels", - save_path=f"labels_{args.dataset}_{args.model}_lr_{args.lr}.png", + save_path=out_path(f"labels_{args.dataset}_{args.model}_lr_{args.lr}.png"), ) if args.task in ["patterns", "mini_patterns"]: @@ -608,7 +696,7 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): p_shapes, labels=shapes_concepts, title="Concepts", - save_path=f"concepts_{args.dataset}_{args.model}_lr_{args.lr}-shapes.png", + save_path=out_path(f"concepts_{args.dataset}_{args.model}_lr_{args.lr}-shapes.png"), ) print("Concept collapse shapes", 1 - compute_coverage(cf_shapes)) @@ -617,7 +705,7 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): p_colors, labels=colors_concepts, title="Concepts", - save_path=f"concepts_{args.dataset}_{args.model}_lr_{args.lr}-colors.png", + save_path=out_path(f"concepts_{args.dataset}_{args.model}_lr_{args.lr}-colors.png"), ) print("Concept collapse colors", 1 - compute_coverage(cf_colors)) @@ -628,7 +716,7 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): c_pred, labels=dataset.get_concept_labels(), title="Concepts", - save_path=f"concepts_{args.dataset}_{args.model}_lr_{args.lr}.png", + save_path=out_path(f"concepts_{args.dataset}_{args.model}_lr_{args.lr}.png"), ) print("Concept collapse", 1 - compute_coverage(cf)) @@ -654,14 +742,32 @@ def train(model: MnistDPL, dataset: BaseDataset, _loss: ADDMNIST_DPL, args): ), } ) - K = max(np.max(c_pred), np.max(c_true)) - wandb.log( - { - "cf-concepts": wandb.plot.confusion_matrix( - None, c_true, c_pred, class_names=[str(i) for i in range(K + 1)] - ), - } - ) + + if args.dataset == "raven": + concept_names, concept_values = dataset.get_concept_labels() + c_true_flat = c_true.reshape(-1, c_true.shape[-1]) + c_pred_flat = c_pred.reshape(-1, c_pred.shape[-1]) + + for i in range(c_pred_flat.shape[1]): + wandb.log( + { + f"cf-concepts-{concept_names[i]}": wandb.plot.confusion_matrix( + None, + c_true_flat[:, i], + c_pred_flat[:, i], + class_names=concept_values[i], + ) + } + ) + else: + K = max(np.max(c_pred), np.max(c_true)) + wandb.log( + { + "cf-concepts": wandb.plot.confusion_matrix( + None, c_true, c_pred, class_names=[str(i) for i in range(K + 1)] + ), + } + ) if hasattr(model, "decoder"): list_images = make_grid(