prototorch_models/examples/gng_iris.py

54 lines
1.2 KiB
Python
Raw Normal View History

"""Growing Neural Gas example using the Iris dataset."""
2021-06-01 15:19:43 +00:00
import argparse
import prototorch as pt
import pytorch_lightning as pl
2021-06-04 20:21:28 +00:00
import torch
2021-06-01 15:19:43 +00:00
if __name__ == "__main__":
# Command-line arguments
parser = argparse.ArgumentParser()
parser = pl.Trainer.add_argparse_args(parser)
args = parser.parse_args()
# Reproducibility
pl.utilities.seed.seed_everything(seed=42)
2021-06-01 15:19:43 +00:00
# Prepare the data
2021-06-04 20:21:28 +00:00
train_ds = pt.datasets.Iris(dims=[0, 2])
2021-06-07 16:35:08 +00:00
train_loader = torch.utils.data.DataLoader(train_ds, batch_size=64)
2021-06-01 15:19:43 +00:00
# Hyperparameters
2021-06-03 13:42:54 +00:00
hparams = dict(
num_prototypes=5,
2021-06-07 16:35:08 +00:00
input_dim=2,
2021-06-03 13:42:54 +00:00
lr=0.1,
)
2021-06-01 15:19:43 +00:00
# Initialize the model
2021-06-04 20:21:28 +00:00
model = pt.models.GrowingNeuralGas(
2021-06-03 13:42:54 +00:00
hparams,
2021-06-04 20:21:28 +00:00
prototype_initializer=pt.components.Zeros(2),
2021-06-03 13:42:54 +00:00
)
2021-06-01 15:19:43 +00:00
2021-06-04 20:21:28 +00:00
# Compute intermediate input and output sizes
model.example_input_array = torch.zeros(4, 2)
2021-06-01 15:19:43 +00:00
# Model summary
print(model)
# Callbacks
vis = pt.models.VisNG2D(data=train_loader)
# Setup trainer
trainer = pl.Trainer.from_argparse_args(
args,
max_epochs=100,
callbacks=[vis],
2021-06-04 20:21:28 +00:00
weights_summary="full",
2021-06-01 15:19:43 +00:00
)
# Training loop
trainer.fit(model, train_loader)