prototorch_models/examples/cbc_iris.py

70 lines
1.8 KiB
Python
Raw Normal View History

2021-04-22 14:01:44 +00:00
"""CBC example using the Iris dataset."""
2021-05-21 15:55:55 +00:00
import argparse
import warnings
2021-05-21 15:55:55 +00:00
import prototorch as pt
2021-04-22 14:01:44 +00:00
import pytorch_lightning as pl
from lightning_fabric.utilities.seed import seed_everything
from prototorch.models import CBC, VisCBC2D
from pytorch_lightning.utilities.warnings import PossibleUserWarning
from torch.utils.data import DataLoader
warnings.filterwarnings("ignore", category=PossibleUserWarning)
warnings.filterwarnings("ignore", category=UserWarning)
2021-04-22 14:01:44 +00:00
if __name__ == "__main__":
# Reproducibility
seed_everything(seed=4)
2021-05-21 15:55:55 +00:00
# Command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument("--gpus", type=int, default=0)
parser.add_argument("--fast_dev_run", type=bool, default=False)
2021-05-21 15:55:55 +00:00
args = parser.parse_args()
2021-04-22 14:01:44 +00:00
# Dataset
train_ds = pt.datasets.Iris(dims=[0, 2])
2021-04-22 14:01:44 +00:00
# Dataloaders
train_loader = DataLoader(train_ds, batch_size=32)
2021-04-22 14:01:44 +00:00
# Hyperparameters
2021-04-23 15:27:47 +00:00
hparams = dict(
distribution=[1, 0, 3],
margin=0.1,
proto_lr=0.01,
bb_lr=0.01,
2021-04-23 15:27:47 +00:00
)
2021-04-22 14:01:44 +00:00
# Initialize the model
model = CBC(
hparams,
components_initializer=pt.initializers.SSCI(train_ds, noise=0.1),
reasonings_initializer=pt.initializers.
PurePositiveReasoningsInitializer(),
)
2021-04-22 14:01:44 +00:00
# Callbacks
vis = VisCBC2D(
data=train_ds,
title="CBC Iris Example",
resolution=100,
axis_off=True,
)
2021-04-22 14:01:44 +00:00
# Setup trainer
trainer = pl.Trainer(
accelerator="cuda" if args.gpus else "cpu",
devices=args.gpus if args.gpus else "auto",
fast_dev_run=args.fast_dev_run,
callbacks=[
vis,
],
detect_anomaly=True,
log_every_n_steps=1,
max_epochs=1000,
2021-04-22 14:01:44 +00:00
)
# Training loop
2021-05-07 13:25:04 +00:00
trainer.fit(model, train_loader)