Add RAVEN benchmark support with RavenDPL (--n_values 3|4)

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 <output>/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
+<output>/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).
+This code is adapted from [Marconato et al. (2024) bears](https://github.com/samuelebortolotti/bears).

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}%")

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)}")