prototorch_models/examples/glvq_mnist.py

120 lines
3.4 KiB
Python
Raw Normal View History

2021-04-21 13:52:42 +00:00
"""GLVQ example using the MNIST dataset.
This script also shows how to use Tensorboard for visualizing the prototypes.
"""
2021-04-21 14:28:20 +00:00
import argparse
2021-04-21 12:54:14 +00:00
import pytorch_lightning as pl
import torchvision
from prototorch.components import initializers as cinit
2021-04-21 13:52:42 +00:00
from torch.utils.data import DataLoader
2021-04-21 12:54:14 +00:00
from torchvision import transforms
from torchvision.datasets import MNIST
2021-04-23 15:27:47 +00:00
from prototorch.models.glvq import ImageGLVQ
2021-04-21 12:54:14 +00:00
2021-04-21 13:52:42 +00:00
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
2021-04-21 12:54:14 +00:00
2021-04-21 13:52:42 +00:00
def on_epoch_end(self, trainer, pl_module):
protos = pl_module.proto_layer.prototypes.detach().cpu()
protos_img = protos.reshape(self.to_shape)
grid = torchvision.utils.make_grid(protos_img, nrow=self.nrow)
# grid = grid.permute((1, 2, 0))
tb = pl_module.logger.experiment
2021-04-23 15:27:47 +00:00
tb.add_image(
tag="MNIST Prototypes",
img_tensor=grid,
global_step=trainer.current_epoch,
dataformats="CHW",
)
2021-04-21 12:54:14 +00:00
if __name__ == "__main__":
2021-04-21 14:28:20 +00:00
# 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()
2021-04-21 13:52:42 +00:00
# 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, ))
]),
)
2021-04-21 12:54:14 +00:00
2021-04-21 13:52:42 +00:00
# Dataloaders
2021-04-21 12:54:14 +00:00
train_loader = DataLoader(mnist_train, batch_size=1024)
2021-04-21 13:52:42 +00:00
test_loader = DataLoader(mnist_test, batch_size=1024)
2021-04-21 12:54:14 +00:00
2021-04-21 13:52:42 +00:00
# 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)
2021-04-21 12:54:14 +00:00
2021-04-23 15:49:29 +00:00
# Hyperparameters
hparams = dict(
2021-04-23 15:27:47 +00:00
input_dim=28 * 28,
nclasses=10,
2021-04-23 15:49:29 +00:00
prototypes_per_class=1,
prototype_initializer=cinit.StratifiedMeanInitializer(x, y),
2021-04-23 15:49:29 +00:00
lr=args.lr,
2021-04-23 15:27:47 +00:00
)
2021-04-23 15:49:29 +00:00
# Initialize the model
model = ImageGLVQ(hparams)
2021-04-23 15:49:29 +00:00
2021-04-21 13:52:42 +00:00
# Model summary
print(model)
2021-04-21 12:54:14 +00:00
2021-04-21 13:52:42 +00:00
# Callbacks
2021-04-21 14:28:20 +00:00
vis = VisualizationCallback(to_shape=(-1, 1, 28, 28), nrow=args.ppc)
2021-04-21 12:54:14 +00:00
2021-04-21 13:52:42 +00:00
# Setup trainer
trainer = pl.Trainer(
2021-04-21 14:28:20 +00:00
gpus=args.gpus, # change to use GPUs for training
max_epochs=args.epochs,
2021-04-21 13:52:42 +00:00
callbacks=[vis],
# accelerator="ddp_cpu", # DEBUG-ONLY
# num_processes=2, # DEBUG-ONLY
)
2021-04-21 12:54:14 +00:00
2021-04-21 13:52:42 +00:00
# Training loop
trainer.fit(model, train_loader, test_loader)