diff --git a/main.py b/main.py index 4c8a8f9..f45d630 100644 --- a/main.py +++ b/main.py @@ -19,13 +19,19 @@ from dataset import (CUB_PRESEGM_PARTS, CUB_PRESEGM_NUM_CLASSES, CUB_PRESEGM_DIR import torch from model import Model import lightning.pytorch as pl +from lightning.fabric.plugins.environments.mpi import MPIEnvironment import os import argparse from utils import set_seed, compute_balanced_accuracy_loader, compute_acc_unordered, get_concept_balance from datetime import datetime from competitors import DNN, CBMDeep, CRM, CMR, _plot_intervenability from torchvision.utils import save_image -from lightning.pytorch.loggers import WandbLogger +from lightning.pytorch.loggers import CSVLogger + +# This workstation has an unrelated mpi4py package but no MPI runtime. The +# released experiment is single-device, so disable Lightning's auto-detection +# probe without altering any scientific code path. +MPIEnvironment.detect = staticmethod(lambda: False) def inspect_val_predictions_with_pause(model, val_loader, concept_names, output_folder, device): @@ -542,10 +548,10 @@ def main(args): if test_only: log_name += "_TESTING" - # Save wandb logs to new_outputs/_/wandb - # You can change the 'entity' argument below to your team name (e.g. entity="my-team-name") + # Local logger: environment-only adaptation, with no change to the model, + # data, loss, optimizer, or evaluation paths. run_name = f"{timestamp}_{dataset}_{args.extra}" if args.extra != "" else f"{timestamp}_{dataset}" - logger = WandbLogger(project="HigherOrderCBMs", entity="conceptlords", save_dir=base_output_folder, name=run_name) + logger = CSVLogger(save_dir=base_output_folder, name=run_name) log_folder = os.path.join(base_output_folder, "checkpoints") # collect all hyperparameters and dataset data for logging @@ -628,6 +634,8 @@ def main(args): if DEVICE == "cpu": trainer = pl.Trainer(accelerator='cpu', max_epochs=epochs, logger=logger, callbacks=selected_callbacks) + elif DEVICE == "mps": + trainer = pl.Trainer(accelerator='mps', devices=1, max_epochs=epochs, logger=logger, callbacks=selected_callbacks) else: trainer = pl.Trainer(accelerator='gpu', devices=[int(DEVICE[-1])], max_epochs=epochs, logger=logger, callbacks=selected_callbacks) if not test_only: @@ -641,6 +649,8 @@ def main(args): print(f"Restoring best model from: {best_model_path}") # pass val_loader and train_loader model = type(model).load_from_checkpoint(best_model_path, val_loader=val_loader, train_loader=train_loader) model.to(DEVICE) + model.ALWAYS_USE_TRUE_MASKS = True if always_use_true_masks else False + model.USE_INITIAL_PROTO_EMBS = True if use_initial_proto_embs else False # print epoch number of best model print(f"Best model epoch: {trainer.current_epoch}") @@ -682,10 +692,8 @@ def main(args): "final_test_c_acc": test_c_acc }) - if isinstance(model, Model) and dataset != "cub": - model.generate_plots(logger=logger, output_folder=output_folder) - - _plot_intervenability(model, logger=logger, nb_runs=1, specific_loader=test_loader, specific_epoch_nb="final_test", is_pgcm=args.model=="PGCM") + # WandB-only plots are omitted; the audit computes and stores the numeric + # prototype/intervention metrics in a separate deterministic driver. if dataset == "celebamask" and args.model == "CBM": inspect_val_predictions_with_pause( diff --git a/model.py b/model.py index 029fb85..5ca81eb 100644 --- a/model.py +++ b/model.py @@ -244,7 +244,8 @@ class Model(pl.LightningModule): if self.reconstruction and not self.swapped and not self.use_pretrained_autoencoder: prototypes_as_features = self.map_proto_to_image(prototype_embeddings) prototype_emb_rec = self.map_proto_image_to_proto(prototypes_as_features) if not self.USE_INITIAL_PROTO_EMBS else prototype_embeddings - prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings) + prototype_embeddings_batch = prototype_embeddings.unsqueeze(0).expand_as(prototype_emb_rec) + prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings_batch) concept_logits_per_proto = self.map_proto_to_concepts(prototype_emb_rec) concept_probs_per_proto = torch.sigmoid(concept_logits_per_proto) prototypes_as_images = prototypes_as_features.unsqueeze(0) if prototypes_as_features.dim() == 2 else prototypes_as_features @@ -325,7 +326,16 @@ class Model(pl.LightningModule): if self.is_cubEMB: return self._forward_cubEMB(x, interventions_mask=interventions_mask, standard_interventions_mask=standard_interventions_mask) - if self.use_pretrained_segmenter: + if self.ALWAYS_USE_TRUE_MASKS and not self.use_pretrained_segmenter: + # Native ColorMNIST supplies exact object masks. This ablation + # isolates the released PGCM/prototype path from the separately + # pretrained segmenter artifact used in the paper environment. + X, M_test, C, Y = x + M = M_test + generated_M = M_test + generated_M_logits = None + segmentation_loss = torch.tensor(0.0, device=self.device) + elif self.use_pretrained_segmenter: segmentation, generated_M, C, Y, X = x generated_M_logits = None M = generated_M @@ -382,7 +392,8 @@ class Model(pl.LightningModule): # prototype_emb_rec = prototype_emb_rec.view(b, nb_proto, -1) # (b, nb_proto, proto_size) # prototype_embeddings_batch = prototype_embeddings.unsqueeze(0).expand(b, -1, -1) # (b, nb_proto, proto_size) # prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings_batch) - prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings) + prototype_embeddings_batch = prototype_embeddings.unsqueeze(0).expand_as(prototype_emb_rec) + prototype_emb_loss = F.mse_loss(prototype_emb_rec, prototype_embeddings_batch) concept_logits_per_proto = self.map_proto_to_concepts(prototype_emb_rec.view(-1, prototype_emb_rec.shape[-1])) # (nb_proto, nb_concepts) concept_probs_per_proto = torch.sigmoid(concept_logits_per_proto) # (nb_proto, nb_concepts) @@ -744,6 +755,12 @@ class Model(pl.LightningModule): prototypes_as_images = prototypes_as_images.unsqueeze(1) # (b, 1, nb_proto, 3, H, W) + target_shape = ( + masked_images.shape[0], masked_images.shape[1], + prototypes_as_images.shape[2], *masked_images.shape[3:] + ) + prototypes_as_images = prototypes_as_images.expand(target_shape) + masked_images = masked_images.expand(target_shape) recons_error = F.mse_loss(prototypes_as_images, masked_images, reduction='none') # (b, nb_objects, nb_proto, 3, H, W) recons_error = recons_error.mean(dim=[3, 4, 5]) # (b, nb_objects, nb_proto) @@ -909,6 +926,12 @@ class Model(pl.LightningModule): prototypes_as_images = prototypes_as_images.unsqueeze(1) # (b, 1, nb_proto, 3, H, W) + target_shape = ( + masked_images.shape[0], masked_images.shape[1], + prototypes_as_images.shape[2], *masked_images.shape[3:] + ) + prototypes_as_images = prototypes_as_images.expand(target_shape) + masked_images = masked_images.expand(target_shape) recons_error = F.mse_loss(prototypes_as_images, masked_images, reduction='none') # (b, nb_objects, nb_proto, 3, H, W) recons_error = recons_error.mean(dim=[3, 4, 5]) # (b, nb_objects, nb_proto) @@ -1056,8 +1079,9 @@ class Model(pl.LightningModule): used_indices = get_used_prototypes_indices(self, self.val_loader) print(f"Used {len(used_indices)} prototypes out of {self.prototypes.weight.shape[0]}") print(f"Used prototypes: {used_indices}") - if not self.is_cubEMB and self.current_epoch % self.plot_frequency == 0: - self.generate_plots(logger=None) + # Visualization is omitted in this local audit because the released + # helper is hard-wired to a WandB logger. Numeric training, + # validation, prototype swapping, and checkpoints are unchanged. if self.current_epoch == self.trainer.max_epochs // 2: diff --git a/neural_networks.py b/neural_networks.py index f74b02c..18350c2 100644 --- a/neural_networks.py +++ b/neural_networks.py @@ -237,7 +237,7 @@ class ResNetUNetSegmenter(nn.Module): # 1. Encoder (Pre-trained ResNet18 or ResNet34 is usually sufficient) # We grab the layers to access intermediate features for skip connections - base_model = models.resnet18(pretrained=True) + base_model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) self.base_layers = list(base_model.children()) self.layer0 = nn.Sequential(*self.base_layers[:3]) # size=(N, 64, x.H/2, x.W/2) @@ -299,6 +299,20 @@ class ResNetUNetSegmenter(nn.Module): class ResNetUNetSegmenterMNIST(nn.Module): def __init__(self, n_class=1): super().__init__() + + # The released forward method references the same four ResNet stages + # and three decoder blocks as the generic segmenter, but the MNIST + # constructor omitted their initialization. Restore those literal + # modules with channel dimensions implied by the released forward. + base_model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) + self.base_layers = list(base_model.children()) + self.layer0 = nn.Sequential(*self.base_layers[:3]) + self.layer1 = nn.Sequential(*self.base_layers[3:5]) + self.layer2 = self.base_layers[5] + self.layer3 = self.base_layers[6] + self.up3 = self._up_block(256, 128) + self.up2 = self._up_block(128 + 128, 64) + self.up1 = self._up_block(64 + 64, 64) self.final_up = nn.Sequential( nn.Conv2d(64, 32, kernel_size=3, padding=1), diff --git a/configs/config_mnist_slot8_cbm.yml b/configs/config_mnist_slot8_cbm.yml new file mode 100644 index 0000000..a5a5e42 --- /dev/null +++ b/configs/config_mnist_slot8_cbm.yml @@ -0,0 +1,27 @@ +dataset: "mnist" +model: "CBM" +epochs: 100 +batch_size: 128 +num_workers: 0 +lr: 0.001 +warmup_epochs: 10 +embedding_size: 128 +rule_emb_size: 100 +concepts_to_task: "thresholding" +use_weights: true +plot_frequency: 1000 +intv_prob: 0.2 +use_balanced_accuracy: true +use_linear_task_predictor: false +segmentation_method: "mask" +use_pretrained_segmenter: false +presegmented_datasets_path: null +use_pretrained_autoencoder: false +fixed_lr: false +always_use_true_masks: true +use_initial_proto_embs: false +noisy_prob: 0.0 +noisy_digit: null +noisy_target_digit: null +noisy_part: null +noisy_target_part: null diff --git a/configs/config_mnist_slot8_cmr.yml b/configs/config_mnist_slot8_cmr.yml new file mode 100644 index 0000000..3b61bf4 --- /dev/null +++ b/configs/config_mnist_slot8_cmr.yml @@ -0,0 +1,15 @@ +dataset: "mnist" +model: "CMR" +epochs: 100 +batch_size: 128 +num_workers: 0 +lr: 0.001 +warmup_epochs: 10 +embedding_size: 128 +rule_emb_size: 100 +concepts_to_task: "thresholding" +use_weights: true +plot_frequency: 1000 +intv_prob: 0.20 +use_balanced_accuracy: true +use_linear_task_predictor: false diff --git a/configs/config_mnist_slot8_crm.yml b/configs/config_mnist_slot8_crm.yml new file mode 100644 index 0000000..d8d1def --- /dev/null +++ b/configs/config_mnist_slot8_crm.yml @@ -0,0 +1,15 @@ +dataset: "mnist" +model: "CRM" +epochs: 100 +batch_size: 128 +num_workers: 0 +lr: 0.001 +warmup_epochs: 10 +embedding_size: 128 +rule_emb_size: 100 +concepts_to_task: "thresholding" +use_weights: true +plot_frequency: 1000 +intv_prob: 0.20 +use_balanced_accuracy: true +use_linear_task_predictor: false diff --git a/configs/config_mnist_slot8_full.yml b/configs/config_mnist_slot8_full.yml new file mode 100644 index 0000000..a29477a --- /dev/null +++ b/configs/config_mnist_slot8_full.yml @@ -0,0 +1,37 @@ +dataset: "mnist" +model: "PGCM" +epochs: 100 +batch_size: 128 +num_workers: 0 +lr: 0.001 +warmup_epochs: 10 +fixed_lr: false +protosize: 64 +embedding_size: 128 +nb_proto: 30 +lam_entropy: 0.01 +lam_batch_entropy: 0.01 +decay_lam_entropy: true +lam_reconstruction: 5.0 +lam_kl: 0.0 +lam_segmentation: 1.0 +lam_orth: 0.0 +lam_proto_emb: 0.0 +segmentation_method: "mask" +plot_frequency: 1000 +concepts_to_task: "thresholding" +use_weights: true +always_use_true_masks: true +use_initial_proto_embs: false +use_balanced_accuracy: true +use_linear_task_predictor: false +use_pretrained_autoencoder: false +autoencoder_path: null +use_pretrained_segmenter: false +presegmented_datasets_path: null +intv_prob: 0.2 +noisy_prob: 0.0 +noisy_digit: null +noisy_target_digit: null +noisy_part: null +noisy_target_part: null diff --git a/configs/config_mnist_slot8_joint_full.yml b/configs/config_mnist_slot8_joint_full.yml new file mode 100644 index 0000000..27f67d2 --- /dev/null +++ b/configs/config_mnist_slot8_joint_full.yml @@ -0,0 +1,37 @@ +dataset: "mnist" +model: "PGCM" +epochs: 100 +batch_size: 128 +num_workers: 0 +lr: 0.001 +warmup_epochs: 10 +fixed_lr: false +protosize: 64 +embedding_size: 128 +nb_proto: 30 +lam_entropy: 0.01 +lam_batch_entropy: 0.01 +decay_lam_entropy: true +lam_reconstruction: 5.0 +lam_kl: 0.0 +lam_segmentation: 1.0 +lam_orth: 0.0 +lam_proto_emb: 0.0 +segmentation_method: "mask" +plot_frequency: 1000 +concepts_to_task: "thresholding" +use_weights: true +always_use_true_masks: false +use_initial_proto_embs: false +use_balanced_accuracy: true +use_linear_task_predictor: false +use_pretrained_autoencoder: false +autoencoder_path: null +use_pretrained_segmenter: false +presegmented_datasets_path: null +intv_prob: 0.2 +noisy_prob: 0.0 +noisy_digit: null +noisy_target_digit: null +noisy_part: null +noisy_target_part: null diff --git a/configs/config_mnist_slot8_joint_noisy.yml b/configs/config_mnist_slot8_joint_noisy.yml new file mode 100644 index 0000000..6d10c33 --- /dev/null +++ b/configs/config_mnist_slot8_joint_noisy.yml @@ -0,0 +1,37 @@ +dataset: "mnist" +model: "PGCM" +epochs: 100 +batch_size: 128 +num_workers: 0 +lr: 0.001 +warmup_epochs: 10 +fixed_lr: false +protosize: 64 +embedding_size: 128 +nb_proto: 30 +lam_entropy: 0.01 +lam_batch_entropy: 0.01 +decay_lam_entropy: true +lam_reconstruction: 5.0 +lam_kl: 0.0 +lam_segmentation: 1.0 +lam_orth: 0.0 +lam_proto_emb: 0.0 +segmentation_method: "mask" +plot_frequency: 1000 +concepts_to_task: "thresholding" +use_weights: true +always_use_true_masks: false +use_initial_proto_embs: false +use_balanced_accuracy: true +use_linear_task_predictor: false +use_pretrained_autoencoder: false +autoencoder_path: null +use_pretrained_segmenter: false +presegmented_datasets_path: null +intv_prob: 0.2 +noisy_prob: 0.3 +noisy_digit: [3, 4] +noisy_target_digit: [1, 8] +noisy_part: null +noisy_target_part: null diff --git a/configs/config_mnist_slot8_joint_smoke.yml b/configs/config_mnist_slot8_joint_smoke.yml new file mode 100644 index 0000000..f87c40f --- /dev/null +++ b/configs/config_mnist_slot8_joint_smoke.yml @@ -0,0 +1,37 @@ +dataset: "mnist" +model: "PGCM" +epochs: 2 +batch_size: 128 +num_workers: 0 +lr: 0.001 +warmup_epochs: 1 +fixed_lr: false +protosize: 64 +embedding_size: 128 +nb_proto: 30 +lam_entropy: 0.01 +lam_batch_entropy: 0.01 +decay_lam_entropy: true +lam_reconstruction: 5.0 +lam_kl: 0.0 +lam_segmentation: 1.0 +lam_orth: 0.0 +lam_proto_emb: 0.0 +segmentation_method: "mask" +plot_frequency: 1000 +concepts_to_task: "thresholding" +use_weights: true +always_use_true_masks: false +use_initial_proto_embs: false +use_balanced_accuracy: true +use_linear_task_predictor: false +use_pretrained_autoencoder: false +autoencoder_path: null +use_pretrained_segmenter: false +presegmented_datasets_path: null +intv_prob: 0.2 +noisy_prob: 0.0 +noisy_digit: null +noisy_target_digit: null +noisy_part: null +noisy_target_part: null diff --git a/configs/config_mnist_slot8_noisy.yml b/configs/config_mnist_slot8_noisy.yml new file mode 100644 index 0000000..800ba31 --- /dev/null +++ b/configs/config_mnist_slot8_noisy.yml @@ -0,0 +1,37 @@ +dataset: "mnist" +model: "PGCM" +epochs: 100 +batch_size: 128 +num_workers: 0 +lr: 0.001 +warmup_epochs: 10 +fixed_lr: false +protosize: 64 +embedding_size: 128 +nb_proto: 30 +lam_entropy: 0.01 +lam_batch_entropy: 0.01 +decay_lam_entropy: true +lam_reconstruction: 5.0 +lam_kl: 0.0 +lam_segmentation: 1.0 +lam_orth: 0.0 +lam_proto_emb: 0.0 +segmentation_method: "mask" +plot_frequency: 1000 +concepts_to_task: "thresholding" +use_weights: true +always_use_true_masks: true +use_initial_proto_embs: false +use_balanced_accuracy: true +use_linear_task_predictor: false +use_pretrained_autoencoder: false +autoencoder_path: null +use_pretrained_segmenter: false +presegmented_datasets_path: null +intv_prob: 0.2 +noisy_prob: 0.3 +noisy_digit: [3, 4] +noisy_target_digit: [1, 8] +noisy_part: null +noisy_target_part: null diff --git a/configs/config_mnist_slot8_smoke.yml b/configs/config_mnist_slot8_smoke.yml new file mode 100644 index 0000000..7ce8c9c --- /dev/null +++ b/configs/config_mnist_slot8_smoke.yml @@ -0,0 +1,37 @@ +dataset: "mnist" +model: "PGCM" +epochs: 2 +batch_size: 128 +num_workers: 0 +lr: 0.001 +warmup_epochs: 1 +fixed_lr: false +protosize: 64 +embedding_size: 128 +nb_proto: 30 +lam_entropy: 0.01 +lam_batch_entropy: 0.01 +decay_lam_entropy: true +lam_reconstruction: 5.0 +lam_kl: 0.0 +lam_segmentation: 1.0 +lam_orth: 0.0 +lam_proto_emb: 0.0 +segmentation_method: "mask" +plot_frequency: 1000 +concepts_to_task: "thresholding" +use_weights: true +always_use_true_masks: true +use_initial_proto_embs: false +use_balanced_accuracy: true +use_linear_task_predictor: false +use_pretrained_autoencoder: false +autoencoder_path: null +use_pretrained_segmenter: false +presegmented_datasets_path: null +intv_prob: 0.2 +noisy_prob: 0.0 +noisy_digit: null +noisy_target_digit: null +noisy_part: null +noisy_target_part: null