update add_distribution_to_annotations to work with dicts and GroupConfig

Author: gdefe

conceptarium/conf/model/_commons.yaml

@@ -17,19 +17,12 @@ latent_encoder_kwargs:
 # Concept distribution configs
 # =============================================================
 variable_distributions:
-  discrete_card1: 
-    path: "torch.distributions.RelaxedBernoulli"
-    kwargs:
-      temperature: 0.1
-  discrete_cardn: 
-    path: "torch.distributions.RelaxedOneHotCategorical"
-    kwargs:
-      temperature: 0.1
-      # num_classes: to be set dynamically for each concept
-  continuous_card1: 
-    path: "torch_concepts.distributions.Delta"
-  continuous_cardn: 
-    path: "torch_concepts.distributions.Delta"
+  _target_: "torch_concepts.GroupConfig"
+  binary: "torch.distributions.RelaxedBernoulli"
+  categorical: "torch.distributions.RelaxedOneHotCategorical"
+  # TODO: handle kwargs
+  # continuous: 
+    # ... not supported yet
 
 
 # =============================================================

examples/utilization/2_model/5_torch_training.py

@@ -13,14 +13,15 @@
 
 import torch
 from torch import nn
-from torch_concepts import Annotations, AxisAnnotation
-from torch_concepts.nn import ConceptBottleneckModel, ConceptLoss
-from torch_concepts.data.datasets import ToyDataset
 from torch.distributions import Bernoulli
 
+from torch_concepts.nn import ConceptBottleneckModel
+from torch_concepts.data.datasets import ToyDataset
+
 from torchmetrics.classification import BinaryAccuracy
 
 
+
 def main():
     # Set random seed for reproducibility
     torch.manual_seed(42)
@@ -33,12 +34,10 @@ def main():
     n_samples = 1000
     dataset = ToyDataset(dataset='xor', seed=42, n_gen=n_samples)
     x_train = dataset.input_data
-    concept_idx = list(dataset.graph.edge_index[0].unique().numpy())
-    task_idx = list(dataset.graph.edge_index[1].unique().numpy())
-    c_train = dataset.concepts[:, concept_idx]
-    y_train = dataset.concepts[:, task_idx]
-    concept_names = [dataset.concept_names[i] for i in concept_idx]
-    task_names = [dataset.concept_names[i] for i in task_idx]
+    c_train = dataset.concepts[:, :2]
+    y_train = dataset.concepts[:, 2:]
+    concept_names = dataset.concept_names[:2]
+    task_names = dataset.concept_names[2:]
 
     n_features = x_train.shape[1]
     n_concepts = c_train.shape[1]
@@ -49,41 +48,25 @@ def main():
     print(f"Tasks: {n_tasks} - {task_names}")
     print(f"Training samples: {n_samples}")
 
-    # For binary concepts, we can use simple labels
-    concept_annotations = Annotations({
-        1: AxisAnnotation(
-            labels=tuple(concept_names + task_names),
-            metadata={
-                concept_names[0]: {
-                    'type': 'discrete',
-                    'distribution': Bernoulli
-                },
-                concept_names[1]: {
-                    'type': 'discrete',
-                    'distribution': Bernoulli
-                },
-                task_names[0]: {
-                    'type': 'discrete',
-                    'distribution': Bernoulli
-                }
-            }
-        )
-    })
+    concept_annotations = dataset.annotations
 
     print(f"Concept axis labels: {concept_annotations[1].labels}")
     print(f"Concept types: {[concept_annotations[1].metadata[name]['type'] for name in concept_names]}")
     print(f"Concept cardinalities: {concept_annotations[1].cardinalities}")
-    print(f"Concept distributions: {[concept_annotations[1].metadata[name]['distribution'] for name in concept_names]}")
 
     # Init model
     print("\n" + "=" * 60)
     print("Step 2: Initialize ConceptBottleneckModel")
     print("=" * 60)
 
+    # Define variable distributions as Bernoulli
+    variable_distributions = {name: Bernoulli for name in concept_names + task_names}
+
     # Initialize the CBM
     model = ConceptBottleneckModel(
         input_size=n_features,
         annotations=concept_annotations,
+        variable_distributions=variable_distributions,
         task_names=task_names,
         latent_encoder_kwargs={'hidden_size': 16, 'n_layers': 1}
     )

examples/utilization/2_model/6_lightning_training.py

@@ -12,32 +12,15 @@
 """
 
 import torch
-from torch.utils.data import Dataset, DataLoader
-from torch_concepts import Annotations, AxisAnnotation
 from torch_concepts.nn import ConceptBottleneckModel
 from torch_concepts.data.datasets import ToyDataset
+from torch_concepts.data.base.datamodule import ConceptDataModule
 from torch.distributions import Bernoulli
 
 from torchmetrics.classification import BinaryAccuracy
 
 from pytorch_lightning import Trainer
 
-class ConceptDataset(Dataset):
-    """Custom dataset that returns batches in the format expected by ConceptBottleneckModel."""
-    
-    def __init__(self, x, c, y):
-        self.x = x
-        self.concepts = torch.cat([c, y], dim=1)
-    
-    def __len__(self):
-        return len(self.x)
-    
-    def __getitem__(self, idx):
-        return {
-            'inputs': {'x': self.x[idx]},
-            'concepts': {'c': self.concepts[idx]},
-        }
-
 def main():
     # Set random seed for reproducibility
     torch.manual_seed(42)
@@ -47,63 +30,39 @@ def main():
     print("Step 1: Generate toy XOR dataset")
     print("=" * 60)
 
-    n_samples = 1000
+    n_samples = 10000
+    batch_size = 2048
     dataset = ToyDataset(dataset='xor', seed=42, n_gen=n_samples)
-    x_train = dataset.input_data
-    concept_idx = list(dataset.graph.edge_index[0].unique().numpy())
-    task_idx = list(dataset.graph.edge_index[1].unique().numpy())
-    c_train = dataset.concepts[:, concept_idx]
-    y_train = dataset.concepts[:, task_idx]
-    concept_names = [dataset.concept_names[i] for i in concept_idx]
-    task_names = [dataset.concept_names[i] for i in task_idx]
-    
-    n_features = x_train.shape[1]
-    n_concepts = c_train.shape[1]
-    n_tasks = y_train.shape[1]
-    
-    print(f"Input features: {n_features}")
-    print(f"Concepts: {n_concepts} - {concept_names}")
-    print(f"Tasks: {n_tasks} - {task_names}")
-    print(f"Training samples: {n_samples}")
+    datamodule = ConceptDataModule(dataset=dataset, 
+                                   batch_size=batch_size,
+                                   val_size=0.1,
+                                   test_size=0.2)
+    annotations = dataset.annotations
+    concept_names = annotations.get_axis_annotation(1).labels
 
-    # For binary concepts, we can use simple labels
-    concept_annotations = Annotations({
-        1: AxisAnnotation(
-            labels=tuple(concept_names + task_names),
-            cardinalities=[1]*(n_concepts + n_tasks),
-            metadata={
-                concept_names[0]: {
-                    'type': 'discrete',
-                    'distribution': Bernoulli
-                },
-                concept_names[1]: {
-                    'type': 'discrete',
-                    'distribution': Bernoulli
-                },
-                task_names[0]: {
-                    'type': 'discrete',
-                    'distribution': Bernoulli
-                }
-            }
-        )
-    })
-    
-    print(f"Concept axis labels: {concept_annotations[1].labels}")
-    print(f"Concept types: {[concept_annotations[1].metadata[name]['type'] for name in concept_names]}")
-    print(f"Concept cardinalities: {concept_annotations[1].cardinalities}")
-    print(f"Concept distributions: {[concept_annotations[1].metadata[name]['distribution'] for name in concept_names]}")
+    n_features = dataset.input_data.shape[1]
+    n_concepts = 2
+    n_tasks = 1
 
+    print(f"Input features: {n_features}")
+    print(f"Concepts: {n_concepts} - {concept_names[:2]}")
+    print(f"Tasks: {n_tasks} - {concept_names[2]}")
+    print(f"Training samples: {n_samples}")
 
     # Init model
     print("\n" + "=" * 60)
     print("Step 2: Initialize ConceptBottleneckModel")
     print("=" * 60)
 
+    # Define variable distributions as Bernoulli
+    variable_distributions = {name: Bernoulli for name in concept_names}
+
     # Initialize the CBM
     model = ConceptBottleneckModel(
         input_size=n_features,
-        annotations=concept_annotations,
-        task_names=task_names,
+        annotations=annotations,
+        variable_distributions=variable_distributions,
+        task_names=['xor'],
         latent_encoder_kwargs={'hidden_size': 16, 'n_layers': 1},
         # Specify loss and optimizer to abilitate training with lightning
         loss=torch.nn.BCEWithLogitsLoss(),
@@ -121,11 +80,10 @@ def main():
     print("Step 3: Test forward pass")
     print("=" * 60)
 
-    batch_size = 8
-    x_batch = x_train[:batch_size]
+    x_batch = dataset.input_data[:batch_size]
 
     # Forward pass
-    query = list(concept_names) + list(task_names)
+    query = concept_names
     print(f"Query variables: {query}")
 
     with torch.no_grad():
@@ -136,43 +94,51 @@ def main():
     print(f"Expected output dim: {n_concepts + n_tasks}")
 
 
-    # Test forward pass
+    # Test lightning training
     print("\n" + "=" * 60)
     print("Step 4: Training loop with lightning")
     print("=" * 60)
 
-    trainer = Trainer(
-        max_epochs=500,
-        log_every_n_steps=10
-    )
-
-    # Create dataset and dataloader
-    train_dataset = ConceptDataset(x_train, c_train, y_train)
-    train_dataloader = DataLoader(train_dataset, batch_size=1000, shuffle=False)
+    trainer = Trainer(max_epochs=100)
 
     model.train()
-    trainer.fit(model, train_dataloaders=train_dataloader)
+    trainer.fit(model, datamodule=datamodule)
 
     # Evaluate
     print("\n" + "=" * 60)
-    print("Step 5: Evaluation")
+    print("Step 5: Evaluation with standard torch metrics")
     print("=" * 60)
 
     concept_acc_fn = BinaryAccuracy()
     task_acc_fn = BinaryAccuracy()
 
     model.eval()
+    concept_acc_sum = 0.0
+    task_acc_sum = 0.0
+    num_batches = 0
+
     with torch.no_grad():
-        endogenous = model(x_train, query=query)
-        c_pred = endogenous[:, :n_concepts]
-        y_pred = endogenous[:, n_concepts:]
-        
-        # Compute accuracy using BinaryAccuracy
-        concept_acc = concept_acc_fn(c_pred, c_train.int()).item()
-        task_acc = task_acc_fn(y_pred, y_train.int()).item()
-        
-        print(f"Concept accuracy: {concept_acc:.4f}")
-        print(f"Task accuracy: {task_acc:.4f}")
+        test_loader = datamodule.test_dataloader()
+        for batch in test_loader:
+            endogenous = model(batch['inputs']['x'], query=query)
+            c_pred = endogenous[:, :n_concepts]
+            y_pred = endogenous[:, n_concepts:]
+
+            c_true = batch['concepts']['c'][:, :n_concepts]
+            y_true = batch['concepts']['c'][:, n_concepts:]
+
+            concept_acc = concept_acc_fn(c_pred, c_true.int()).item()
+            task_acc = task_acc_fn(y_pred, y_true.int()).item()
+
+            concept_acc_sum += concept_acc
+            task_acc_sum += task_acc
+            num_batches += 1
+
+    avg_concept_acc = concept_acc_sum / num_batches if num_batches > 0 else 0.0
+    avg_task_acc = task_acc_sum / num_batches if num_batches > 0 else 0.0
+
+    print(f"Average concept accuracy: {avg_concept_acc:.4f}")
+    print(f"Average task accuracy: {avg_task_acc:.4f}")
 
 if __name__ == "__main__":
     main()