Add more CBC examples. MNIST is broken.

examples/cbc_circle.py

@@ -0,0 +1,121 @@
+"""CBC example using the Iris dataset."""
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from matplotlib import pyplot as plt
+from prototorch.models.cbc import CBC, rescaled_cosine_similarity, euclidean_similarity
+from prototorch.models.glvq import GLVQ
+from sklearn.datasets import make_circles
+from torch.utils.data import DataLoader, TensorDataset
+
+
+class NumpyDataset(TensorDataset):
+    def __init__(self, *arrays):
+        # tensors = [torch.from_numpy(arr) for arr in arrays]
+        tensors = [torch.Tensor(arr) for arr in arrays]
+        super().__init__(*tensors)
+
+
+class VisualizationCallback(pl.Callback):
+    def __init__(self,
+                 x_train,
+                 y_train,
+                 prototype_model=True,
+                 title="Prototype Visualization",
+                 cmap="viridis"):
+        super().__init__()
+        self.x_train = x_train
+        self.y_train = y_train
+        self.title = title
+        self.fig = plt.figure(self.title)
+        self.cmap = cmap
+        self.prototype_model = prototype_model
+
+    def on_epoch_end(self, trainer, pl_module):
+        if self.prototype_model:
+            protos = pl_module.prototypes
+            color = pl_module.prototype_labels
+        else:
+            protos = pl_module.components
+            color = 'k'
+        ax = self.fig.gca()
+        ax.cla()
+        ax.set_title(self.title)
+        ax.set_xlabel("Data dimension 1")
+        ax.set_ylabel("Data dimension 2")
+        ax.scatter(x_train[:, 0], x_train[:, 1], c=y_train, edgecolor="k")
+        ax.scatter(protos[:, 0],
+                   protos[:, 1],
+                   c=color,
+                   cmap=self.cmap,
+                   edgecolor="k",
+                   marker="D",
+                   s=50)
+        x = np.vstack((x_train, protos))
+        x_min, x_max = x[:, 0].min() - 1, x[:, 0].max() + 1
+        y_min, y_max = x[:, 1].min() - 1, x[:, 1].max() + 1
+        xx, yy = np.meshgrid(np.arange(x_min, x_max, 1 / 50),
+                             np.arange(y_min, y_max, 1 / 50))
+        mesh_input = np.c_[xx.ravel(), yy.ravel()]
+        y_pred = pl_module.predict(torch.Tensor(mesh_input))
+        y_pred = y_pred.reshape(xx.shape)
+
+        ax.contourf(xx, yy, y_pred, cmap=self.cmap, alpha=0.35)
+        ax.set_xlim(left=x_min + 0, right=x_max - 0)
+        ax.set_ylim(bottom=y_min + 0, top=y_max - 0)
+        plt.pause(0.1)
+
+
+if __name__ == "__main__":
+    # Dataset
+    x_train, y_train = make_circles(n_samples=300,
+                                    shuffle=True,
+                                    noise=0.05,
+                                    random_state=None,
+                                    factor=0.5)
+    train_ds = NumpyDataset(x_train, y_train)
+
+    # Dataloaders
+    train_loader = DataLoader(train_ds, num_workers=0, batch_size=150)
+
+    # Hyperparameters
+    hparams = dict(
+        input_dim=x_train.shape[1],
+        nclasses=len(np.unique(y_train)),
+        prototypes_per_class=5,
+        prototype_initializer="randn",
+        lr=0.01,
+    )
+
+    # Initialize the model
+    model = CBC(
+        hparams,
+        data=[x_train, y_train],
+        similarity=euclidean_similarity,
+    )
+
+    #model = GLVQ(hparams, data=[x_train, y_train])
+
+    # Fix the component locations
+    # model.proto_layer.requires_grad_(False)
+
+    # import sys
+    # sys.exit()
+
+    # Model summary
+    print(model)
+
+    # Callbacks
+    vis = VisualizationCallback(x_train, y_train, prototype_model=False)
+
+    # Setup trainer
+    trainer = pl.Trainer(
+        max_epochs=500,
+        callbacks=[
+            vis,
+        ],
+    )
+
+    # Training loop
+    trainer.fit(model, train_loader)

examples/cbc_mnist.py

@@ -0,0 +1,128 @@
+"""CBC example using the MNIST dataset.
+
+This script also shows how to use Tensorboard for visualizing the prototypes.
+"""
+
+import argparse
+
+import pytorch_lightning as pl
+import torchvision
+from matplotlib import pyplot as plt
+from prototorch.models.cbc import ImageCBC, euclidean_similarity, rescaled_cosine_similarity
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from torchvision.datasets import MNIST
+
+
+class VisualizationCallback(pl.Callback):
+    def __init__(self, to_shape=(-1, 1, 28, 28), nrow=2):
+        super().__init__()
+        self.to_shape = to_shape
+        self.nrow = nrow
+
+    def on_epoch_end(self, trainer, pl_module: ImageCBC):
+        tb = pl_module.logger.experiment
+
+        # components
+        components = pl_module.components
+        components_img = components.reshape(self.to_shape)
+        grid = torchvision.utils.make_grid(components_img, nrow=self.nrow)
+        tb.add_image(
+            tag="MNIST Components",
+            img_tensor=grid,
+            global_step=trainer.current_epoch,
+            dataformats="CHW",
+        )
+        # Reasonings
+        reasonings = pl_module.reasonings
+        tb.add_images(
+            tag="MNIST Reasoning",
+            img_tensor=reasonings,
+            global_step=trainer.current_epoch,
+            dataformats="NCHW",
+        )
+
+
+if __name__ == "__main__":
+    # Arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--epochs",
+                        type=int,
+                        default=10,
+                        help="Epochs to train.")
+    parser.add_argument("--lr",
+                        type=float,
+                        default=0.001,
+                        help="Learning rate.")
+    parser.add_argument("--batch_size",
+                        type=int,
+                        default=256,
+                        help="Batch size.")
+    parser.add_argument("--gpus",
+                        type=int,
+                        default=0,
+                        help="Number of GPUs to use.")
+    parser.add_argument("--ppc",
+                        type=int,
+                        default=1,
+                        help="Prototypes-Per-Class.")
+    args = parser.parse_args()
+
+    # Dataset
+    mnist_train = MNIST(
+        "./datasets",
+        train=True,
+        download=True,
+        transform=transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.1307, ), (0.3081, ))
+        ]),
+    )
+    mnist_test = MNIST(
+        "./datasets",
+        train=False,
+        download=True,
+        transform=transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.1307, ), (0.3081, ))
+        ]),
+    )
+
+    # Dataloaders
+    train_loader = DataLoader(mnist_train, batch_size=1024)
+    test_loader = DataLoader(mnist_test, batch_size=1024)
+
+    # Grab the full dataset to warm-start prototypes
+    x, y = next(iter(DataLoader(mnist_train, batch_size=len(mnist_train))))
+    x = x.view(len(mnist_train), -1)
+
+    # Hyperparameters
+    hparams = dict(
+        input_dim=28 * 28,
+        nclasses=10,
+        prototypes_per_class=args.ppc,
+        prototype_initializer="randn",
+        lr=1,
+        similarity=euclidean_similarity,
+    )
+
+    # Initialize the model
+    model = ImageCBC(hparams, data=[x, y])
+    # Model summary
+    print(model)
+
+    # Callbacks
+    vis = VisualizationCallback(to_shape=(-1, 1, 28, 28), nrow=args.ppc)
+
+    # Setup trainer
+    trainer = pl.Trainer(
+        gpus=args.gpus,  # change to use GPUs for training
+        max_epochs=args.epochs,
+        callbacks=[vis],
+        track_grad_norm=2,
+        # accelerator="ddp_cpu",  # DEBUG-ONLY
+        # num_processes=2,  # DEBUG-ONLY
+    )
+
+    # Training loop
+    trainer.fit(model, train_loader, test_loader)

prototorch/models/cbc.py

@@ -33,14 +33,12 @@ class CosineSimilarity(torch.nn.Module):
 
     def forward(self, x, y):
         epsilon = torch.finfo(x.dtype).eps
-        normed_x = (x/ x.pow(2) \
+        normed_x = (x / x.pow(2).sum(dim=tuple(range(
-            .sum(dim=tuple(range(1, x.ndim)), keepdim=True) \
+            1, x.ndim)), keepdim=True).clamp(min=epsilon).sqrt()).flatten(
-            .clamp(min=epsilon) \
+                start_dim=1)
-            .sqrt()).flatten(start_dim=1)
+        normed_y = (y / y.pow(2).sum(dim=tuple(range(
-        normed_y = (y / y.pow(2) \
+            1, y.ndim)), keepdim=True).clamp(min=epsilon).sqrt()).flatten(
-            .sum(dim=tuple(range(1, y.ndim)), keepdim=True) \
+                start_dim=1)
-            .clamp(min=epsilon) \
-            .sqrt()).flatten(start_dim=1)
         # normed_x = (x / torch.linalg.norm(x, dim=1))
         diss = torch.inner(normed_x, normed_y)
         return self.activation(diss)
@@ -73,14 +71,14 @@ class ReasoningLayer(torch.nn.Module):
         probabilities_init.uniform_(0.4, 0.6)
         self.reasoning_probabilities = torch.nn.Parameter(probabilities_init)
 
-    # @property
+    @property
-    # def reasonings(self):
+    def reasonings(self):
-    #     pk = self.reasoning_probabilities[0]
+        pk = self.reasoning_probabilities[0]
-    #     nk = (1 - pk) * self.reasoning_probabilities[1]
+        nk = (1 - pk) * self.reasoning_probabilities[1]
-    #     ik = (1 - pk - nk)
+        ik = (1 - pk - nk)
-    #     # pk is of shape (1, n_components, n_classes)
+        # pk is of shape (1, n_components, n_classes)
-    #     img = torch.cat([pk, nk, ik], dim=0).permute(1, 0, 2)
+        img = torch.cat([pk, nk, ik], dim=0).permute(1, 0, 2)
-    #     return img.unsqueeze(1)  # (n_components, 1, 3, n_classes)
+        return img.unsqueeze(1)  # (n_components, 1, 3, n_classes)
 
     def forward(self, detections):
         pk = self.reasoning_probabilities[0].clamp(0, 1)
@@ -128,7 +126,11 @@ class CBC(pl.LightningModule):
 
     @property
     def components(self):
-        return self.proto_layer.prototypes.detach().numpy()
+        return self.proto_layer.prototypes.detach().cpu()
+
+    @property
+    def reasonings(self):
+        return self.reasoning_layer.reasonings.cpu()
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
@@ -140,8 +142,8 @@ class CBC(pl.LightningModule):
 
     def forward(self, x):
         self.sync_backbones()
-        # protos = self.proto_layer.prototypes
+        protos = self.proto_layer.prototypes
-        protos, _ = self.proto_layer()
+        # protos, _ = self.proto_layer()
 
         latent_x = self.backbone(x)
         latent_protos = self.backbone_dependent(protos)
@@ -163,7 +165,7 @@ class CBC(pl.LightningModule):
         # y_true = torch.nn.functional.one_hot(y, num_classes=nclasses)
         # y_true = torch.eye(nclasses)[y.long()]
         y_true = torch.nn.functional.one_hot(y.long(), num_classes=nclasses)
-        loss = MarginLoss(self.margin)(y_pred, y_true).sum(dim=0)
+        loss = MarginLoss(self.margin)(y_pred, y_true).mean(dim=0)
         self.log("train_loss", loss)
         # with torch.no_grad():
         #     preds = torch.argmax(y_pred, dim=1)
@@ -172,12 +174,14 @@ class CBC(pl.LightningModule):
         #     preds.int(),
         #     y.int())  # FloatTensors are assumed to be class probabilities
         self.train_acc(y_pred, y_true)
-        self.log("acc",
+        self.log(
+            "acc",
             self.train_acc,
             on_step=False,
             on_epoch=True,
             prog_bar=True,
-                 logger=True)
+            logger=True,
+        )
         return loss
 
     #def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
@@ -201,5 +205,5 @@ class ImageCBC(CBC):
     clamping after updates.
     """
     def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
-        super().on_train_batch_end(outputs, batch, batch_idx, dataload_idx)
+        #super().on_train_batch_end(outputs, batch, batch_idx, dataloader_idx)
-        self.proto_layer.prototypes.data.clamp_(0., 1.)
+        self.proto_layer.prototypes.data.clamp_(0.0, 1.0)