julius/prototorch_models
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add more CBC examples. MNIST is broken.

Browse Source
This commit is contained in:
Alexander Engelsberger 2021-04-22 17:37:20 +02:00
parent 2e2f6707f6
commit db4499a103
3 changed files with 283 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

121
examples/cbc_circle.py Normal file
View File

@ -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)

128
examples/cbc_mnist.py Normal file
View File

@ -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)

52
prototorch/models/cbc.py
View File

@ -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) \
.sum(dim=tuple(range(1, x.ndim)), keepdim=True) \
.clamp(min=epsilon) \
.sqrt()).flatten(start_dim=1)
normed_y = (y / y.pow(2) \
.sum(dim=tuple(range(1, y.ndim)), keepdim=True) \
.clamp(min=epsilon) \
.sqrt()).flatten(start_dim=1)
normed_x = (x / x.pow(2).sum(dim=tuple(range(
1, x.ndim)), keepdim=True).clamp(min=epsilon).sqrt()).flatten(
start_dim=1)
normed_y = (y / y.pow(2).sum(dim=tuple(range(
1, y.ndim)), keepdim=True).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
# def reasonings(self):
# pk = self.reasoning_probabilities[0]
# nk = (1 - pk) * self.reasoning_probabilities[1]
# ik = (1 - pk - nk)
# # pk is of shape (1, n_components, n_classes)
# img = torch.cat([pk, nk, ik], dim=0).permute(1, 0, 2)
# return img.unsqueeze(1) # (n_components, 1, 3, n_classes)
@property
def reasonings(self):
pk = self.reasoning_probabilities[0]
nk = (1 - pk) * self.reasoning_probabilities[1]
ik = (1 - pk - nk)
# pk is of shape (1, n_components, n_classes)
img = torch.cat([pk, nk, ik], dim=0).permute(1, 0, 2)
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()
protos = self.proto_layer.prototypes
# 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)
self.proto_layer.prototypes.data.clamp_(0., 1.)
#super().on_train_batch_end(outputs, batch, batch_idx, dataloader_idx)
self.proto_layer.prototypes.data.clamp_(0.0, 1.0)