137 lines
3.4 KiB
Python
137 lines
3.4 KiB
Python
import logging
|
|
|
|
import pytorch_lightning as pl
|
|
import torchmetrics
|
|
from prototorch.core import SMCI, PCALinearTransformInitializer
|
|
from prototorch.datasets import Iris
|
|
from prototorch.models.architectures.base import Steps
|
|
from prototorch.models.callbacks import (
|
|
LogTorchmetricCallback,
|
|
PlotLambdaMatrixToTensorboard,
|
|
VisGMLVQ2D,
|
|
)
|
|
from prototorch.models.library.gmlvq import GMLVQ
|
|
from pytorch_lightning.callbacks import EarlyStopping
|
|
from torch.utils.data import DataLoader, random_split
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
# ##############################################################################
|
|
|
|
|
|
def main():
|
|
# ------------------------------------------------------------
|
|
# DATA
|
|
# ------------------------------------------------------------
|
|
|
|
# Dataset
|
|
full_dataset = Iris()
|
|
full_count = len(full_dataset)
|
|
|
|
train_count = int(full_count * 0.5)
|
|
val_count = int(full_count * 0.4)
|
|
test_count = int(full_count * 0.1)
|
|
|
|
train_dataset, val_dataset, test_dataset = random_split(
|
|
full_dataset, (train_count, val_count, test_count))
|
|
|
|
# Dataloader
|
|
train_loader = DataLoader(
|
|
train_dataset,
|
|
batch_size=1,
|
|
num_workers=4,
|
|
shuffle=True,
|
|
)
|
|
val_loader = DataLoader(
|
|
val_dataset,
|
|
batch_size=1,
|
|
num_workers=4,
|
|
shuffle=False,
|
|
)
|
|
test_loader = DataLoader(
|
|
test_dataset,
|
|
batch_size=1,
|
|
num_workers=0,
|
|
shuffle=False,
|
|
)
|
|
|
|
# ------------------------------------------------------------
|
|
# HYPERPARAMETERS
|
|
# ------------------------------------------------------------
|
|
|
|
# Select Initializer
|
|
components_initializer = SMCI(full_dataset)
|
|
|
|
# Define Hyperparameters
|
|
hyperparameters = GMLVQ.HyperParameters(
|
|
lr=dict(components_layer=0.1, _omega=0),
|
|
input_dim=4,
|
|
distribution=dict(
|
|
num_classes=3,
|
|
per_class=1,
|
|
),
|
|
component_initializer=components_initializer,
|
|
omega_initializer=PCALinearTransformInitializer,
|
|
omega_initializer_kwargs=dict(
|
|
data=train_dataset.dataset[train_dataset.indices][0]))
|
|
|
|
# Create Model
|
|
model = GMLVQ(hyperparameters)
|
|
|
|
# ------------------------------------------------------------
|
|
# TRAINING
|
|
# ------------------------------------------------------------
|
|
|
|
# Controlling Callbacks
|
|
recall = LogTorchmetricCallback(
|
|
'training_recall',
|
|
torchmetrics.Recall,
|
|
num_classes=3,
|
|
step=Steps.TRAINING,
|
|
)
|
|
|
|
stopping_criterion = LogTorchmetricCallback(
|
|
'validation_recall',
|
|
torchmetrics.Recall,
|
|
num_classes=3,
|
|
step=Steps.VALIDATION,
|
|
)
|
|
|
|
es = EarlyStopping(
|
|
monitor=stopping_criterion.name,
|
|
mode="max",
|
|
patience=10,
|
|
)
|
|
|
|
# Visualization Callback
|
|
vis = VisGMLVQ2D(data=full_dataset)
|
|
|
|
# Define trainer
|
|
trainer = pl.Trainer(
|
|
callbacks=[
|
|
vis,
|
|
recall,
|
|
stopping_criterion,
|
|
es,
|
|
PlotLambdaMatrixToTensorboard(),
|
|
],
|
|
max_epochs=100,
|
|
)
|
|
|
|
# Train
|
|
trainer.fit(model, train_loader, val_loader)
|
|
trainer.test(model, test_loader)
|
|
|
|
# Manual save
|
|
trainer.save_checkpoint("./y_arch.ckpt")
|
|
|
|
# Load saved model
|
|
new_model = GMLVQ.load_from_checkpoint(
|
|
checkpoint_path="./y_arch.ckpt",
|
|
strict=True,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|