build: bump version 0.5.1 → 0.5.2

fix(hotfix): Protobuf error workaround
build: bump version 0.5.0 → 0.5.1
2022-06-01 14:25:44 +02:00 · 2022-06-01 14:14:57 +02:00 · 2022-05-17 12:04:08 +02:00 · 2022-05-17 12:03:43 +02:00 · 2022-05-17 12:00:52 +02:00 · 2022-05-16 11:12:53 +02:00
34 changed files with 690 additions and 359 deletions
--- a/.bumpversion.cfg
+++ b/.bumpversion.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.5.0
+current_version = 0.5.2
 commit = True
 tag = True
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -3,7 +3,7 @@

 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.1.0
+  rev: v4.2.0
  hooks:
  - id: trailing-whitespace
  - id: end-of-file-fixer
@@ -23,7 +23,7 @@ repos:
  - id: isort

 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.931
+  rev: v0.950
  hooks:
  - id: mypy
    files: prototorch
@@ -42,7 +42,7 @@ repos:
  - id: python-check-blanket-noqa

 - repo: https://github.com/asottile/pyupgrade
-  rev: v2.31.0
+  rev: v2.32.1
  hooks:
  - id: pyupgrade

--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -23,7 +23,7 @@ author = "Jensun Ravichandran"

 # The full version, including alpha/beta/rc tags
 #
-release = "0.5.0"
+release = "0.5.2"

 # -- General configuration ---------------------------------------------------

--- a/examples/cbc_iris.py
+++ b/examples/cbc_iris.py
@@ -1,12 +1,22 @@
 """CBC example using the Iris dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
-import torch
+from prototorch.models import CBC, VisCBC2D
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
+    # Reproducibility
+    seed_everything(seed=4)
+
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -15,11 +25,8 @@ if __name__ == "__main__":
    # Dataset
    train_ds = pt.datasets.Iris(dims=[0, 2])

-    # Reproducibility
-    pl.utilities.seed.seed_everything(seed=42)
-
    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=32)
+    train_loader = DataLoader(train_ds, batch_size=32)

    # Hyperparameters
    hparams = dict(
@@ -30,23 +37,30 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.CBC(
+    model = CBC(
        hparams,
-        components_initializer=pt.initializers.SSCI(train_ds, noise=0.01),
-        reasonings_iniitializer=pt.initializers.
+        components_initializer=pt.initializers.SSCI(train_ds, noise=0.1),
+        reasonings_initializer=pt.initializers.
        PurePositiveReasoningsInitializer(),
    )

    # Callbacks
-    vis = pt.models.VisCBC2D(data=train_ds,
-                             title="CBC Iris Example",
-                             resolution=100,
-                             axis_off=True)
+    vis = VisCBC2D(
+        data=train_ds,
+        title="CBC Iris Example",
+        resolution=100,
+        axis_off=True,
+    )

    # Setup trainer
    trainer = pl.Trainer.from_argparse_args(
        args,
-        callbacks=[vis],
+        callbacks=[
+            vis,
+        ],
+        detect_anomaly=True,
+        log_every_n_steps=1,
+        max_epochs=1000,
    )

    # Training loop
--- a/examples/dynamic_pruning.py
+++ b/examples/dynamic_pruning.py
@@ -1,12 +1,29 @@
 """Dynamically prune 'loser' prototypes in GLVQ-type models."""

 import argparse
+import logging
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import (
+    CELVQ,
+    PruneLoserPrototypes,
+    VisGLVQ2D,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
+    # Reproducibility
+    seed_everything(seed=4)
+
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -16,15 +33,17 @@ if __name__ == "__main__":
    num_classes = 4
    num_features = 2
    num_clusters = 1
-    train_ds = pt.datasets.Random(num_samples=500,
-                                  num_classes=num_classes,
-                                  num_features=num_features,
-                                  num_clusters=num_clusters,
-                                  separation=3.0,
-                                  seed=42)
+    train_ds = pt.datasets.Random(
+        num_samples=500,
+        num_classes=num_classes,
+        num_features=num_features,
+        num_clusters=num_clusters,
+        separation=3.0,
+        seed=42,
+    )

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=256)
+    train_loader = DataLoader(train_ds, batch_size=256)

    # Hyperparameters
    prototypes_per_class = num_clusters * 5
@@ -34,7 +53,7 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.CELVQ(
+    model = CELVQ(
        hparams,
        prototypes_initializer=pt.initializers.FVCI(2, 3.0),
    )
@@ -43,18 +62,18 @@ if __name__ == "__main__":
    model.example_input_array = torch.zeros(4, 2)

    # Summary
-    print(model)
+    logging.info(model)

    # Callbacks
-    vis = pt.models.VisGLVQ2D(train_ds)
-    pruning = pt.models.PruneLoserPrototypes(
+    vis = VisGLVQ2D(train_ds)
+    pruning = PruneLoserPrototypes(
        threshold=0.01,  # prune prototype if it wins less than 1%
        idle_epochs=20,  # pruning too early may cause problems
        prune_quota_per_epoch=2,  # prune at most 2 prototypes per epoch
        frequency=1,  # prune every epoch
        verbose=True,
    )
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
        monitor="train_loss",
        min_delta=0.001,
        patience=20,
@@ -71,10 +90,9 @@ if __name__ == "__main__":
            pruning,
            es,
        ],
-        progress_bar_refresh_rate=0,
-        terminate_on_nan=True,
-        weights_summary="full",
-        accelerator="ddp",
+        detect_anomaly=True,
+        log_every_n_steps=1,
+        max_epochs=1000,
    )

    # Training loop
--- a/examples/glvq_iris.py
+++ b/examples/glvq_iris.py
@@ -1,13 +1,24 @@
 """GLVQ example using the Iris dataset."""

 import argparse
+import logging
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import GLVQ, VisGLVQ2D
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
 from torch.optim.lr_scheduler import ExponentialLR
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=PossibleUserWarning)

 if __name__ == "__main__":
+    # Reproducibility
+    seed_everything(seed=4)
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -17,7 +28,7 @@ if __name__ == "__main__":
    train_ds = pt.datasets.Iris(dims=[0, 2])

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=64)
+    train_loader = DataLoader(train_ds, batch_size=64, num_workers=4)

    # Hyperparameters
    hparams = dict(
@@ -29,7 +40,7 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.GLVQ(
+    model = GLVQ(
        hparams,
        optimizer=torch.optim.Adam,
        prototypes_initializer=pt.initializers.SMCI(train_ds),
@@ -41,14 +52,17 @@ if __name__ == "__main__":
    model.example_input_array = torch.zeros(4, 2)

    # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
+    vis = VisGLVQ2D(data=train_ds)

    # Setup trainer
    trainer = pl.Trainer.from_argparse_args(
        args,
-        callbacks=[vis],
-        weights_summary="full",
-        accelerator="ddp",
+        callbacks=[
+            vis,
+        ],
+        max_epochs=100,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
@@ -58,8 +72,8 @@ if __name__ == "__main__":
    trainer.save_checkpoint("./glvq_iris.ckpt")

    # Load saved model
-    new_model = pt.models.GLVQ.load_from_checkpoint(
+    new_model = GLVQ.load_from_checkpoint(
        checkpoint_path="./glvq_iris.ckpt",
        strict=False,
    )
-    print(new_model)
+    logging.info(new_model)
--- a/examples/gmlvq_iris.py
+++ b/examples/gmlvq_iris.py
@@ -1,13 +1,25 @@
 """GMLVQ example using the Iris dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import GMLVQ, VisGMLVQ2D
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
 from torch.optim.lr_scheduler import ExponentialLR
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
+
+    # Reproducibility
+    seed_everything(seed=4)
+
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -17,7 +29,7 @@ if __name__ == "__main__":
    train_ds = pt.datasets.Iris()

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=64)
+    train_loader = DataLoader(train_ds, batch_size=64)

    # Hyperparameters
    hparams = dict(
@@ -32,7 +44,7 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.GMLVQ(
+    model = GMLVQ(
        hparams,
        optimizer=torch.optim.Adam,
        prototypes_initializer=pt.initializers.SMCI(train_ds),
@@ -44,14 +56,17 @@ if __name__ == "__main__":
    model.example_input_array = torch.zeros(4, 4)

    # Callbacks
-    vis = pt.models.VisGMLVQ2D(data=train_ds)
+    vis = VisGMLVQ2D(data=train_ds)

    # Setup trainer
    trainer = pl.Trainer.from_argparse_args(
        args,
-        callbacks=[vis],
-        weights_summary="full",
-        accelerator="ddp",
+        callbacks=[
+            vis,
+        ],
+        max_epochs=100,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/gmlvq_mnist.py
+++ b/examples/gmlvq_mnist.py
@@ -1,14 +1,29 @@
 """GMLVQ example using the MNIST dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import (
+    ImageGMLVQ,
+    PruneLoserPrototypes,
+    VisImgComp,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
 from torchvision import transforms
 from torchvision.datasets import MNIST

+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)
+
 if __name__ == "__main__":
+    # Reproducibility
+    seed_everything(seed=4)
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -33,12 +48,8 @@ if __name__ == "__main__":
    )

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds,
-                                               num_workers=0,
-                                               batch_size=256)
-    test_loader = torch.utils.data.DataLoader(test_ds,
-                                              num_workers=0,
-                                              batch_size=256)
+    train_loader = DataLoader(train_ds, num_workers=4, batch_size=256)
+    test_loader = DataLoader(test_ds, num_workers=4, batch_size=256)

    # Hyperparameters
    num_classes = 10
@@ -52,14 +63,14 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.ImageGMLVQ(
+    model = ImageGMLVQ(
        hparams,
        optimizer=torch.optim.Adam,
        prototypes_initializer=pt.initializers.SMCI(train_ds),
    )

    # Callbacks
-    vis = pt.models.VisImgComp(
+    vis = VisImgComp(
        data=train_ds,
        num_columns=10,
        show=False,
@@ -69,14 +80,14 @@ if __name__ == "__main__":
        embedding_data=200,
        flatten_data=False,
    )
-    pruning = pt.models.PruneLoserPrototypes(
+    pruning = PruneLoserPrototypes(
        threshold=0.01,
        idle_epochs=1,
        prune_quota_per_epoch=10,
        frequency=1,
        verbose=True,
    )
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
        monitor="train_loss",
        min_delta=0.001,
        patience=15,
@@ -90,11 +101,11 @@ if __name__ == "__main__":
        callbacks=[
            vis,
            pruning,
-            # es,
+            es,
        ],
-        terminate_on_nan=True,
-        weights_summary=None,
-        # accelerator="ddp",
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/gmlvq_spiral.py
+++ b/examples/gmlvq_spiral.py
@@ -1,12 +1,28 @@
 """GMLVQ example using the spiral dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import (
+    GMLVQ,
+    PruneLoserPrototypes,
+    VisGLVQ2D,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
+    # Reproducibility
+    seed_everything(seed=4)
+
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -16,7 +32,7 @@ if __name__ == "__main__":
    train_ds = pt.datasets.Spiral(num_samples=500, noise=0.5)

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=256)
+    train_loader = DataLoader(train_ds, batch_size=256)

    # Hyperparameters
    num_classes = 2
@@ -32,19 +48,19 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.GMLVQ(
+    model = GMLVQ(
        hparams,
        optimizer=torch.optim.Adam,
        prototypes_initializer=pt.initializers.SSCI(train_ds, noise=1e-2),
    )

    # Callbacks
-    vis = pt.models.VisGLVQ2D(
+    vis = VisGLVQ2D(
        train_ds,
        show_last_only=False,
        block=False,
    )
-    pruning = pt.models.PruneLoserPrototypes(
+    pruning = PruneLoserPrototypes(
        threshold=0.01,
        idle_epochs=10,
        prune_quota_per_epoch=5,
@@ -53,7 +69,7 @@ if __name__ == "__main__":
        prototypes_initializer=pt.initializers.SSCI(train_ds, noise=1e-1),
        verbose=True,
    )
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
        monitor="train_loss",
        min_delta=1.0,
        patience=5,
@@ -69,7 +85,9 @@ if __name__ == "__main__":
            es,
            pruning,
        ],
-        terminate_on_nan=True,
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/gng_iris.py
+++ b/examples/gng_iris.py
@@ -1,10 +1,19 @@
 """Growing Neural Gas example using the Iris dataset."""

 import argparse
+import logging
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import GrowingNeuralGas, VisNG2D
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
    # Command-line arguments
@@ -13,11 +22,11 @@ if __name__ == "__main__":
    args = parser.parse_args()

    # Reproducibility
-    pl.utilities.seed.seed_everything(seed=42)
+    seed_everything(seed=42)

    # Prepare the data
    train_ds = pt.datasets.Iris(dims=[0, 2])
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=64)
+    train_loader = DataLoader(train_ds, batch_size=64)

    # Hyperparameters
    hparams = dict(
@@ -27,7 +36,7 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.GrowingNeuralGas(
+    model = GrowingNeuralGas(
        hparams,
        prototypes_initializer=pt.initializers.ZCI(2),
    )
@@ -36,17 +45,20 @@ if __name__ == "__main__":
    model.example_input_array = torch.zeros(4, 2)

    # Model summary
-    print(model)
+    logging.info(model)

    # Callbacks
-    vis = pt.models.VisNG2D(data=train_loader)
+    vis = VisNG2D(data=train_loader)

    # Setup trainer
    trainer = pl.Trainer.from_argparse_args(
        args,
+        callbacks=[
+            vis,
+        ],
        max_epochs=100,
-        callbacks=[vis],
-        weights_summary="full",
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/gtlvq_mnist.py
+++ b/examples/gtlvq_mnist.py
@@ -1,14 +1,30 @@
 """GTLVQ example using the MNIST dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import (
+    ImageGTLVQ,
+    PruneLoserPrototypes,
+    VisImgComp,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
 from torchvision import transforms
 from torchvision.datasets import MNIST

+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)
+
 if __name__ == "__main__":
+    # Reproducibility
+    seed_everything(seed=4)
+
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -33,12 +49,8 @@ if __name__ == "__main__":
    )

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds,
-                                               num_workers=0,
-                                               batch_size=256)
-    test_loader = torch.utils.data.DataLoader(test_ds,
-                                              num_workers=0,
-                                              batch_size=256)
+    train_loader = DataLoader(train_ds, num_workers=0, batch_size=256)
+    test_loader = DataLoader(test_ds, num_workers=0, batch_size=256)

    # Hyperparameters
    num_classes = 10
@@ -52,7 +64,7 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.ImageGTLVQ(
+    model = ImageGTLVQ(
        hparams,
        optimizer=torch.optim.Adam,
        prototypes_initializer=pt.initializers.SMCI(train_ds),
@@ -61,7 +73,7 @@ if __name__ == "__main__":
            next(iter(train_loader))[0].reshape(256, 28 * 28)))

    # Callbacks
-    vis = pt.models.VisImgComp(
+    vis = VisImgComp(
        data=train_ds,
        num_columns=10,
        show=False,
@@ -71,14 +83,14 @@ if __name__ == "__main__":
        embedding_data=200,
        flatten_data=False,
    )
-    pruning = pt.models.PruneLoserPrototypes(
+    pruning = PruneLoserPrototypes(
        threshold=0.01,
        idle_epochs=1,
        prune_quota_per_epoch=10,
        frequency=1,
        verbose=True,
    )
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
        monitor="train_loss",
        min_delta=0.001,
        patience=15,
@@ -93,11 +105,11 @@ if __name__ == "__main__":
        callbacks=[
            vis,
            pruning,
-            # es,
+            es,
        ],
-        terminate_on_nan=True,
-        weights_summary=None,
-        accelerator="ddp",
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/gtlvq_moons.py
+++ b/examples/gtlvq_moons.py
@@ -1,10 +1,20 @@
 """Localized-GTLVQ example using the Moons dataset."""

 import argparse
+import logging
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import GTLVQ, VisGLVQ2D
+from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
    # Command-line arguments
@@ -13,33 +23,35 @@ if __name__ == "__main__":
    args = parser.parse_args()

    # Reproducibility
-    pl.utilities.seed.seed_everything(seed=2)
+    seed_everything(seed=2)

    # Dataset
    train_ds = pt.datasets.Moons(num_samples=300, noise=0.2, seed=42)

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds,
-                                               batch_size=256,
-                                               shuffle=True)
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=256,
+        shuffle=True,
+    )

    # Hyperparameters
    # Latent_dim should be lower than input dim.
    hparams = dict(distribution=[1, 3], input_dim=2, latent_dim=1)

    # Initialize the model
-    model = pt.models.GTLVQ(
-        hparams, prototypes_initializer=pt.initializers.SMCI(train_ds))
+    model = GTLVQ(hparams,
+                  prototypes_initializer=pt.initializers.SMCI(train_ds))

    # Compute intermediate input and output sizes
    model.example_input_array = torch.zeros(4, 2)

    # Summary
-    print(model)
+    logging.info(model)

    # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
-    es = pl.callbacks.EarlyStopping(
+    vis = VisGLVQ2D(data=train_ds)
+    es = EarlyStopping(
        monitor="train_acc",
        min_delta=0.001,
        patience=20,
@@ -55,8 +67,9 @@ if __name__ == "__main__":
            vis,
            es,
        ],
-        weights_summary="full",
-        accelerator="ddp",
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/knn_iris.py
+++ b/examples/knn_iris.py
@@ -1,12 +1,19 @@
 """k-NN example using the Iris dataset from scikit-learn."""

 import argparse
+import logging
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import KNN, VisGLVQ2D
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
 from sklearn.datasets import load_iris
 from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)

 if __name__ == "__main__":
    # Command-line arguments
@@ -16,34 +23,36 @@ if __name__ == "__main__":

    # Dataset
    X, y = load_iris(return_X_y=True)
-    X = X[:, [0, 2]]
+    X = X[:, 0:3:2]

-    X_train, X_test, y_train, y_test = train_test_split(X,
-                                                        y,
-                                                        test_size=0.5,
-                                                        random_state=42)
+    X_train, X_test, y_train, y_test = train_test_split(
+        X,
+        y,
+        test_size=0.5,
+        random_state=42,
+    )

    train_ds = pt.datasets.NumpyDataset(X_train, y_train)
    test_ds = pt.datasets.NumpyDataset(X_test, y_test)

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=16)
-    test_loader = torch.utils.data.DataLoader(test_ds, batch_size=16)
+    train_loader = DataLoader(train_ds, batch_size=16)
+    test_loader = DataLoader(test_ds, batch_size=16)

    # Hyperparameters
    hparams = dict(k=5)

    # Initialize the model
-    model = pt.models.KNN(hparams, data=train_ds)
+    model = KNN(hparams, data=train_ds)

    # Compute intermediate input and output sizes
    model.example_input_array = torch.zeros(4, 2)

    # Summary
-    print(model)
+    logging.info(model)

    # Callbacks
-    vis = pt.models.VisGLVQ2D(
+    vis = VisGLVQ2D(
        data=(X_train, y_train),
        resolution=200,
        block=True,
@@ -53,8 +62,11 @@ if __name__ == "__main__":
    trainer = pl.Trainer.from_argparse_args(
        args,
        max_epochs=1,
-        callbacks=[vis],
-        weights_summary="full",
+        callbacks=[
+            vis,
+        ],
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
@@ -63,7 +75,7 @@ if __name__ == "__main__":

    # Recall
    y_pred = model.predict(torch.tensor(X_train))
-    print(y_pred)
+    logging.info(y_pred)

    # Test
    trainer.test(model, dataloaders=test_loader)
--- a/examples/ksom_colors.py
+++ b/examples/ksom_colors.py
@@ -1,12 +1,21 @@
 """Kohonen Self Organizing Map."""

 import argparse
+import logging
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
 from matplotlib import pyplot as plt
+from prototorch.models import KohonenSOM
 from prototorch.utils.colors import hex_to_rgb
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader, TensorDataset
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)


 class Vis2DColorSOM(pl.Callback):
@@ -18,7 +27,7 @@ class Vis2DColorSOM(pl.Callback):
        self.data = data
        self.pause_time = pause_time

-    def on_epoch_end(self, trainer, pl_module):
+    def on_train_epoch_end(self, trainer, pl_module: KohonenSOM):
        ax = self.fig.gca()
        ax.cla()
        ax.set_title(self.title)
@@ -31,12 +40,14 @@ class Vis2DColorSOM(pl.Callback):
        d = pl_module.compute_distances(self.data)
        wp = pl_module.predict_from_distances(d)
        for i, iloc in enumerate(wp):
-            plt.text(iloc[1],
-                     iloc[0],
-                     cnames[i],
-                     ha="center",
-                     va="center",
-                     bbox=dict(facecolor="white", alpha=0.5, lw=0))
+            plt.text(
+                iloc[1],
+                iloc[0],
+                color_names[i],
+                ha="center",
+                va="center",
+                bbox=dict(facecolor="white", alpha=0.5, lw=0),
+            )

        if trainer.current_epoch != trainer.max_epochs - 1:
            plt.pause(self.pause_time)
@@ -51,7 +62,7 @@ if __name__ == "__main__":
    args = parser.parse_args()

    # Reproducibility
-    pl.utilities.seed.seed_everything(seed=42)
+    seed_everything(seed=42)

    # Prepare the data
    hex_colors = [
@@ -59,15 +70,15 @@ if __name__ == "__main__":
        "#00ff00", "#ff0000", "#00ffff", "#ff00ff", "#ffff00", "#ffffff",
        "#545454", "#7f7f7f", "#a8a8a8", "#808000", "#800080", "#ffa500"
    ]
-    cnames = [
+    color_names = [
        "black", "blue", "darkblue", "skyblue", "greyblue", "lilac", "green",
        "red", "cyan", "magenta", "yellow", "white", "darkgrey", "mediumgrey",
        "lightgrey", "olive", "purple", "orange"
    ]
    colors = list(hex_to_rgb(hex_colors))
    data = torch.Tensor(colors) / 255.0
-    train_ds = torch.utils.data.TensorDataset(data)
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=8)
+    train_ds = TensorDataset(data)
+    train_loader = DataLoader(train_ds, batch_size=8)

    # Hyperparameters
    hparams = dict(
@@ -78,7 +89,7 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.KohonenSOM(
+    model = KohonenSOM(
        hparams,
        prototypes_initializer=pt.initializers.RNCI(3),
    )
@@ -87,7 +98,7 @@ if __name__ == "__main__":
    model.example_input_array = torch.zeros(4, 3)

    # Model summary
-    print(model)
+    logging.info(model)

    # Callbacks
    vis = Vis2DColorSOM(data=data)
@@ -96,8 +107,11 @@ if __name__ == "__main__":
    trainer = pl.Trainer.from_argparse_args(
        args,
        max_epochs=500,
-        callbacks=[vis],
-        weights_summary="full",
+        callbacks=[
+            vis,
+        ],
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/lgmlvq_moons.py
+++ b/examples/lgmlvq_moons.py
@@ -1,10 +1,20 @@
 """Localized-GMLVQ example using the Moons dataset."""

 import argparse
+import logging
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import LGMLVQ, VisGLVQ2D
+from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
    # Command-line arguments
@@ -13,15 +23,13 @@ if __name__ == "__main__":
    args = parser.parse_args()

    # Reproducibility
-    pl.utilities.seed.seed_everything(seed=2)
+    seed_everything(seed=2)

    # Dataset
    train_ds = pt.datasets.Moons(num_samples=300, noise=0.2, seed=42)

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds,
-                                               batch_size=256,
-                                               shuffle=True)
+    train_loader = DataLoader(train_ds, batch_size=256, shuffle=True)

    # Hyperparameters
    hparams = dict(
@@ -31,7 +39,7 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.LGMLVQ(
+    model = LGMLVQ(
        hparams,
        prototypes_initializer=pt.initializers.SMCI(train_ds),
    )
@@ -40,11 +48,11 @@ if __name__ == "__main__":
    model.example_input_array = torch.zeros(4, 2)

    # Summary
-    print(model)
+    logging.info(model)

    # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
-    es = pl.callbacks.EarlyStopping(
+    vis = VisGLVQ2D(data=train_ds)
+    es = EarlyStopping(
        monitor="train_acc",
        min_delta=0.001,
        patience=20,
@@ -60,8 +68,9 @@ if __name__ == "__main__":
            vis,
            es,
        ],
-        weights_summary="full",
-        accelerator="ddp",
+        log_every_n_steps=1,
+        max_epochs=1000,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/lvqmln_iris.py
+++ b/examples/lvqmln_iris.py
@@ -1,10 +1,22 @@
 """LVQMLN example using all four dimensions of the Iris dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import (
+    LVQMLN,
+    PruneLoserPrototypes,
+    VisSiameseGLVQ2D,
+)
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)


 class Backbone(torch.nn.Module):
@@ -34,10 +46,10 @@ if __name__ == "__main__":
    train_ds = pt.datasets.Iris()

    # Reproducibility
-    pl.utilities.seed.seed_everything(seed=42)
+    seed_everything(seed=42)

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=150)
+    train_loader = DataLoader(train_ds, batch_size=150)

    # Hyperparameters
    hparams = dict(
@@ -50,7 +62,7 @@ if __name__ == "__main__":
    backbone = Backbone()

    # Initialize the model
-    model = pt.models.LVQMLN(
+    model = LVQMLN(
        hparams,
        prototypes_initializer=pt.initializers.SSCI(
            train_ds,
@@ -59,18 +71,15 @@ if __name__ == "__main__":
        backbone=backbone,
    )

-    # Model summary
-    print(model)
-
    # Callbacks
-    vis = pt.models.VisSiameseGLVQ2D(
+    vis = VisSiameseGLVQ2D(
        data=train_ds,
        map_protos=False,
        border=0.1,
        resolution=500,
        axis_off=True,
    )
-    pruning = pt.models.PruneLoserPrototypes(
+    pruning = PruneLoserPrototypes(
        threshold=0.01,
        idle_epochs=20,
        prune_quota_per_epoch=2,
@@ -85,6 +94,9 @@ if __name__ == "__main__":
            vis,
            pruning,
        ],
+        log_every_n_steps=1,
+        max_epochs=1000,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/median_lvq_iris.py
+++ b/examples/median_lvq_iris.py
@@ -1,12 +1,23 @@
 """Median-LVQ example using the Iris dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import MedianLVQ, VisGLVQ2D
+from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
+    # Reproducibility
+    seed_everything(seed=4)
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -16,13 +27,13 @@ if __name__ == "__main__":
    train_ds = pt.datasets.Iris(dims=[0, 2])

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(
+    train_loader = DataLoader(
        train_ds,
        batch_size=len(train_ds),  # MedianLVQ cannot handle mini-batches
    )

    # Initialize the model
-    model = pt.models.MedianLVQ(
+    model = MedianLVQ(
        hparams=dict(distribution=(3, 2), lr=0.01),
        prototypes_initializer=pt.initializers.SSCI(train_ds),
    )
@@ -31,8 +42,8 @@ if __name__ == "__main__":
    model.example_input_array = torch.zeros(4, 2)

    # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
-    es = pl.callbacks.EarlyStopping(
+    vis = VisGLVQ2D(data=train_ds)
+    es = EarlyStopping(
        monitor="train_acc",
        min_delta=0.01,
        patience=5,
@@ -44,8 +55,13 @@ if __name__ == "__main__":
    # Setup trainer
    trainer = pl.Trainer.from_argparse_args(
        args,
-        callbacks=[vis, es],
-        weights_summary="full",
+        callbacks=[
+            vis,
+            es,
+        ],
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/ng_iris.py
+++ b/examples/ng_iris.py
@@ -1,15 +1,26 @@
 """Neural Gas example using the Iris dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import NeuralGas, VisNG2D
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
 from sklearn.datasets import load_iris
 from sklearn.preprocessing import StandardScaler
 from torch.optim.lr_scheduler import ExponentialLR
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
+    # Reproducibility
+    seed_everything(seed=4)
+
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -17,7 +28,7 @@ if __name__ == "__main__":

    # Prepare and pre-process the dataset
    x_train, y_train = load_iris(return_X_y=True)
-    x_train = x_train[:, [0, 2]]
+    x_train = x_train[:, 0:3:2]
    scaler = StandardScaler()
    scaler.fit(x_train)
    x_train = scaler.transform(x_train)
@@ -25,7 +36,7 @@ if __name__ == "__main__":
    train_ds = pt.datasets.NumpyDataset(x_train, y_train)

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=150)
+    train_loader = DataLoader(train_ds, batch_size=150)

    # Hyperparameters
    hparams = dict(
@@ -35,7 +46,7 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.NeuralGas(
+    model = NeuralGas(
        hparams,
        prototypes_initializer=pt.core.ZCI(2),
        lr_scheduler=ExponentialLR,
@@ -45,17 +56,18 @@ if __name__ == "__main__":
    # Compute intermediate input and output sizes
    model.example_input_array = torch.zeros(4, 2)

-    # Model summary
-    print(model)
-
    # Callbacks
-    vis = pt.models.VisNG2D(data=train_ds)
+    vis = VisNG2D(data=train_ds)

    # Setup trainer
    trainer = pl.Trainer.from_argparse_args(
        args,
-        callbacks=[vis],
-        weights_summary="full",
+        callbacks=[
+            vis,
+        ],
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/rslvq_iris.py
+++ b/examples/rslvq_iris.py
@@ -1,10 +1,18 @@
 """RSLVQ example using the Iris dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import RSLVQ, VisGLVQ2D
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)

 if __name__ == "__main__":
    # Command-line arguments
@@ -13,13 +21,13 @@ if __name__ == "__main__":
    args = parser.parse_args()

    # Reproducibility
-    pl.utilities.seed.seed_everything(seed=42)
+    seed_everything(seed=42)

    # Dataset
    train_ds = pt.datasets.Iris(dims=[0, 2])

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=64)
+    train_loader = DataLoader(train_ds, batch_size=64)

    # Hyperparameters
    hparams = dict(
@@ -33,7 +41,7 @@ if __name__ == "__main__":
    )

    # Initialize the model
-    model = pt.models.RSLVQ(
+    model = RSLVQ(
        hparams,
        optimizer=torch.optim.Adam,
        prototypes_initializer=pt.initializers.SSCI(train_ds, noise=0.2),
@@ -42,19 +50,18 @@ if __name__ == "__main__":
    # Compute intermediate input and output sizes
    model.example_input_array = torch.zeros(4, 2)

-    # Summary
-    print(model)
-
    # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
+    vis = VisGLVQ2D(data=train_ds)

    # Setup trainer
    trainer = pl.Trainer.from_argparse_args(
        args,
-        callbacks=[vis],
-        terminate_on_nan=True,
-        weights_summary="full",
-        accelerator="ddp",
+        callbacks=[
+            vis,
+        ],
+        detect_anomaly=True,
+        max_epochs=100,
+        log_every_n_steps=1,
    )

    # Training loop
--- a/examples/siamese_glvq_iris.py
+++ b/examples/siamese_glvq_iris.py
@@ -1,10 +1,18 @@
 """Siamese GLVQ example using all four dimensions of the Iris dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import SiameseGLVQ, VisSiameseGLVQ2D
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)


 class Backbone(torch.nn.Module):
@@ -34,10 +42,10 @@ if __name__ == "__main__":
    train_ds = pt.datasets.Iris()

    # Reproducibility
-    pl.utilities.seed.seed_everything(seed=2)
+    seed_everything(seed=2)

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=150)
+    train_loader = DataLoader(train_ds, batch_size=150)

    # Hyperparameters
    hparams = dict(
@@ -50,23 +58,25 @@ if __name__ == "__main__":
    backbone = Backbone()

    # Initialize the model
-    model = pt.models.SiameseGLVQ(
+    model = SiameseGLVQ(
        hparams,
        prototypes_initializer=pt.initializers.SMCI(train_ds),
        backbone=backbone,
        both_path_gradients=False,
    )

-    # Model summary
-    print(model)
-
    # Callbacks
-    vis = pt.models.VisSiameseGLVQ2D(data=train_ds, border=0.1)
+    vis = VisSiameseGLVQ2D(data=train_ds, border=0.1)

    # Setup trainer
    trainer = pl.Trainer.from_argparse_args(
        args,
-        callbacks=[vis],
+        callbacks=[
+            vis,
+        ],
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/siamese_gtlvq_iris.py
+++ b/examples/siamese_gtlvq_iris.py
@@ -1,10 +1,18 @@
 """Siamese GTLVQ example using all four dimensions of the Iris dataset."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import SiameseGTLVQ, VisSiameseGLVQ2D
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)
+warnings.filterwarnings("ignore", category=UserWarning)


 class Backbone(torch.nn.Module):
@@ -34,39 +42,43 @@ if __name__ == "__main__":
    train_ds = pt.datasets.Iris()

    # Reproducibility
-    pl.utilities.seed.seed_everything(seed=2)
+    seed_everything(seed=2)

    # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=150)
+    train_loader = DataLoader(train_ds, batch_size=150)

    # Hyperparameters
-    hparams = dict(distribution=[1, 2, 3],
-                   proto_lr=0.01,
-                   bb_lr=0.01,
-                   input_dim=2,
-                   latent_dim=1)
+    hparams = dict(
+        distribution=[1, 2, 3],
+        proto_lr=0.01,
+        bb_lr=0.01,
+        input_dim=2,
+        latent_dim=1,
+    )

    # Initialize the backbone
    backbone = Backbone(latent_size=hparams["input_dim"])

    # Initialize the model
-    model = pt.models.SiameseGTLVQ(
+    model = SiameseGTLVQ(
        hparams,
        prototypes_initializer=pt.initializers.SMCI(train_ds),
        backbone=backbone,
        both_path_gradients=False,
    )

-    # Model summary
-    print(model)
-
    # Callbacks
-    vis = pt.models.VisSiameseGLVQ2D(data=train_ds, border=0.1)
+    vis = VisSiameseGLVQ2D(data=train_ds, border=0.1)

    # Setup trainer
    trainer = pl.Trainer.from_argparse_args(
        args,
-        callbacks=[vis],
+        callbacks=[
+            vis,
+        ],
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/examples/warm_starting.py
+++ b/examples/warm_starting.py
@@ -1,13 +1,30 @@
 """Warm-starting GLVQ with prototypes from Growing Neural Gas."""

 import argparse
+import warnings

 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from prototorch.models import (
+    GLVQ,
+    KNN,
+    GrowingNeuralGas,
+    PruneLoserPrototypes,
+    VisGLVQ2D,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.utilities.seed import seed_everything
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
 from torch.optim.lr_scheduler import ExponentialLR
+from torch.utils.data import DataLoader
+
+warnings.filterwarnings("ignore", category=PossibleUserWarning)

 if __name__ == "__main__":
+
+    # Reproducibility
+    seed_everything(seed=4)
    # Command-line arguments
    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
@@ -15,10 +32,10 @@ if __name__ == "__main__":

    # Prepare the data
    train_ds = pt.datasets.Iris(dims=[0, 2])
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=64)
+    train_loader = DataLoader(train_ds, batch_size=64, num_workers=0)

    # Initialize the gng
-    gng = pt.models.GrowingNeuralGas(
+    gng = GrowingNeuralGas(
        hparams=dict(num_prototypes=5, insert_freq=2, lr=0.1),
        prototypes_initializer=pt.initializers.ZCI(2),
        lr_scheduler=ExponentialLR,
@@ -26,7 +43,7 @@ if __name__ == "__main__":
    )

    # Callbacks
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
        monitor="loss",
        min_delta=0.001,
        patience=20,
@@ -37,9 +54,12 @@ if __name__ == "__main__":

    # Setup trainer for GNG
    trainer = pl.Trainer(
-        max_epochs=100,
-        callbacks=[es],
-        weights_summary=None,
+        max_epochs=1000,
+        callbacks=[
+            es,
+        ],
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
@@ -52,12 +72,12 @@ if __name__ == "__main__":
    )

    # Warm-start prototypes
-    knn = pt.models.KNN(dict(k=1), data=train_ds)
+    knn = KNN(dict(k=1), data=train_ds)
    prototypes = gng.prototypes
    plabels = knn.predict(prototypes)

    # Initialize the model
-    model = pt.models.GLVQ(
+    model = GLVQ(
        hparams,
        optimizer=torch.optim.Adam,
        prototypes_initializer=pt.initializers.LCI(prototypes),
@@ -70,15 +90,15 @@ if __name__ == "__main__":
    model.example_input_array = torch.zeros(4, 2)

    # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
-    pruning = pt.models.PruneLoserPrototypes(
+    vis = VisGLVQ2D(data=train_ds)
+    pruning = PruneLoserPrototypes(
        threshold=0.02,
        idle_epochs=2,
        prune_quota_per_epoch=5,
        frequency=1,
        verbose=True,
    )
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
        monitor="train_loss",
        min_delta=0.001,
        patience=10,
@@ -95,8 +115,9 @@ if __name__ == "__main__":
            pruning,
            es,
        ],
-        weights_summary="full",
-        accelerator="ddp",
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
    )

    # Training loop
--- a/prototorch/models/init.py
+++ b/prototorch/models/init.py
@@ -36,4 +36,4 @@ from .unsupervised import (
 )
 from .vis import *

-__version__ = "0.5.0"
+__version__ = "0.5.2"
--- a/prototorch/models/abstract.py
+++ b/prototorch/models/abstract.py
@@ -1,15 +1,24 @@
 """Abstract classes to be inherited by prototorch models."""

+import logging
+
 import pytorch_lightning as pl
 import torch
+import torch.nn.functional as F
 import torchmetrics
-
-from ..core.competitions import WTAC
-from ..core.components import Components, LabeledComponents
-from ..core.distances import euclidean_distance
-from ..core.initializers import LabelsInitializer, ZerosCompInitializer
-from ..core.pooling import stratified_min_pooling
-from ..nn.wrappers import LambdaLayer
+from prototorch.core.competitions import WTAC
+from prototorch.core.components import (
+    AbstractComponents,
+    Components,
+    LabeledComponents,
+)
+from prototorch.core.distances import euclidean_distance
+from prototorch.core.initializers import (
+    LabelsInitializer,
+    ZerosCompInitializer,
+)
+from prototorch.core.pooling import stratified_min_pooling
+from prototorch.nn.wrappers import LambdaLayer


 class ProtoTorchBolt(pl.LightningModule):
@@ -30,7 +39,7 @@ class ProtoTorchBolt(pl.LightningModule):
        self.lr_scheduler_kwargs = kwargs.get("lr_scheduler_kwargs", dict())

    def configure_optimizers(self):
-        optimizer = self.optimizer(self.parameters(), lr=self.hparams.lr)
+        optimizer = self.optimizer(self.parameters(), lr=self.hparams["lr"])
        if self.lr_scheduler is not None:
            scheduler = self.lr_scheduler(optimizer,
                                          **self.lr_scheduler_kwargs)
@@ -43,7 +52,10 @@ class ProtoTorchBolt(pl.LightningModule):
            return optimizer

    def reconfigure_optimizers(self):
-        self.trainer.strategy.setup_optimizers(self.trainer)
+        if self.trainer:
+            self.trainer.strategy.setup_optimizers(self.trainer)
+        else:
+            logging.warning("No trainer to reconfigure optimizers!")

    def __repr__(self):
        surep = super().__repr__()
@@ -53,6 +65,7 @@ class ProtoTorchBolt(pl.LightningModule):


 class PrototypeModel(ProtoTorchBolt):
+    proto_layer: AbstractComponents

    def __init__(self, hparams, **kwargs):
        super().__init__(hparams, **kwargs)
@@ -75,16 +88,17 @@ class PrototypeModel(ProtoTorchBolt):

    def add_prototypes(self, *args, **kwargs):
        self.proto_layer.add_components(*args, **kwargs)
-        self.hparams.distribution = self.proto_layer.distribution
+        self.hparams["distribution"] = self.proto_layer.distribution
        self.reconfigure_optimizers()

    def remove_prototypes(self, indices):
        self.proto_layer.remove_components(indices)
-        self.hparams.distribution = self.proto_layer.distribution
+        self.hparams["distribution"] = self.proto_layer.distribution
        self.reconfigure_optimizers()


 class UnsupervisedPrototypeModel(PrototypeModel):
+    proto_layer: Components

    def __init__(self, hparams, **kwargs):
        super().__init__(hparams, **kwargs)
@@ -93,7 +107,7 @@ class UnsupervisedPrototypeModel(PrototypeModel):
        prototypes_initializer = kwargs.get("prototypes_initializer", None)
        if prototypes_initializer is not None:
            self.proto_layer = Components(
-                self.hparams.num_prototypes,
+                self.hparams["num_prototypes"],
                initializer=prototypes_initializer,
            )

@@ -108,6 +122,7 @@ class UnsupervisedPrototypeModel(PrototypeModel):


 class SupervisedPrototypeModel(PrototypeModel):
+    proto_layer: LabeledComponents

    def __init__(self, hparams, skip_proto_layer=False, **kwargs):
        super().__init__(hparams, **kwargs)
@@ -127,13 +142,13 @@ class SupervisedPrototypeModel(PrototypeModel):
                    labels_initializer=labels_initializer,
                )
                proto_shape = self.proto_layer.components.shape[1:]
-                self.hparams.initialized_proto_shape = proto_shape
+                self.hparams["initialized_proto_shape"] = proto_shape
            else:
                # when restoring a checkpointed model
                self.proto_layer = LabeledComponents(
                    distribution=distribution,
                    components_initializer=ZerosCompInitializer(
-                        self.hparams.initialized_proto_shape),
+                        self.hparams["initialized_proto_shape"]),
                )
        self.competition_layer = WTAC()

@@ -154,7 +169,7 @@ class SupervisedPrototypeModel(PrototypeModel):
        distances = self.compute_distances(x)
        _, plabels = self.proto_layer()
        winning = stratified_min_pooling(distances, plabels)
-        y_pred = torch.nn.functional.softmin(winning, dim=1)
+        y_pred = F.softmin(winning, dim=1)
        return y_pred

    def predict_from_distances(self, distances):
@@ -207,8 +222,10 @@ class NonGradientMixin(ProtoTorchMixin):

 class ImagePrototypesMixin(ProtoTorchMixin):
    """Mixin for models with image prototypes."""
+    proto_layer: Components
+    components: torch.Tensor

-    def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
+    def on_train_batch_end(self, outputs, batch, batch_idx):
        """Constrain the components to the range [0, 1] by clamping after updates."""
        self.proto_layer.components.data.clamp_(0.0, 1.0)

--- a/prototorch/models/callbacks.py
+++ b/prototorch/models/callbacks.py
@@ -1,25 +1,30 @@
 """Lightning Callbacks."""

 import logging
+from typing import TYPE_CHECKING

 import pytorch_lightning as pl
 import torch
+from prototorch.core.initializers import LiteralCompInitializer

-from ..core.components import Components
-from ..core.initializers import LiteralCompInitializer
 from .extras import ConnectionTopology

+if TYPE_CHECKING:
+    from prototorch.models import GLVQ, GrowingNeuralGas
+

 class PruneLoserPrototypes(pl.Callback):

-    def __init__(self,
-                 threshold=0.01,
-                 idle_epochs=10,
-                 prune_quota_per_epoch=-1,
-                 frequency=1,
-                 replace=False,
-                 prototypes_initializer=None,
-                 verbose=False):
+    def __init__(
+        self,
+        threshold=0.01,
+        idle_epochs=10,
+        prune_quota_per_epoch=-1,
+        frequency=1,
+        replace=False,
+        prototypes_initializer=None,
+        verbose=False,
+    ):
        self.threshold = threshold  # minimum win ratio
        self.idle_epochs = idle_epochs  # epochs to wait before pruning
        self.prune_quota_per_epoch = prune_quota_per_epoch
@@ -28,7 +33,7 @@ class PruneLoserPrototypes(pl.Callback):
        self.verbose = verbose
        self.prototypes_initializer = prototypes_initializer

-    def on_epoch_end(self, trainer, pl_module):
+    def on_train_epoch_end(self, trainer, pl_module: "GLVQ"):
        if (trainer.current_epoch + 1) < self.idle_epochs:
            return None
        if (trainer.current_epoch + 1) % self.frequency:
@@ -43,27 +48,29 @@ class PruneLoserPrototypes(pl.Callback):
            prune_labels = prune_labels[:self.prune_quota_per_epoch]

        if len(to_prune) > 0:
-            if self.verbose:
-                print(f"\nPrototype win ratios: {ratios}")
-                print(f"Pruning prototypes at: {to_prune}")
-                print(f"Corresponding labels are: {prune_labels.tolist()}")
+            logging.debug(f"\nPrototype win ratios: {ratios}")
+            logging.debug(f"Pruning prototypes at: {to_prune}")
+            logging.debug(f"Corresponding labels are: {prune_labels.tolist()}")
+
            cur_num_protos = pl_module.num_prototypes
            pl_module.remove_prototypes(indices=to_prune)
+
            if self.replace:
                labels, counts = torch.unique(prune_labels,
                                              sorted=True,
                                              return_counts=True)
                distribution = dict(zip(labels.tolist(), counts.tolist()))
-                if self.verbose:
-                    print(f"Re-adding pruned prototypes...")
-                    print(f"distribution={distribution}")
+
+                logging.info(f"Re-adding pruned prototypes...")
+                logging.debug(f"distribution={distribution}")
+
                pl_module.add_prototypes(
                    distribution=distribution,
                    components_initializer=self.prototypes_initializer)
            new_num_protos = pl_module.num_prototypes
-            if self.verbose:
-                print(f"`num_prototypes` changed from {cur_num_protos} "
-                      f"to {new_num_protos}.")
+
+            logging.info(f"`num_prototypes` changed from {cur_num_protos} "
+                         f"to {new_num_protos}.")
        return True


@@ -74,11 +81,11 @@ class PrototypeConvergence(pl.Callback):
        self.idle_epochs = idle_epochs  # epochs to wait
        self.verbose = verbose

-    def on_epoch_end(self, trainer, pl_module):
+    def on_train_epoch_end(self, trainer, pl_module):
        if (trainer.current_epoch + 1) < self.idle_epochs:
            return None
-        if self.verbose:
-            print("Stopping...")
+
+        logging.info("Stopping...")
        # TODO
        return True

@@ -96,12 +103,16 @@ class GNGCallback(pl.Callback):
        self.reduction = reduction
        self.freq = freq

-    def on_epoch_end(self, trainer: pl.Trainer, pl_module):
+    def on_train_epoch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: "GrowingNeuralGas",
+    ):
        if (trainer.current_epoch + 1) % self.freq == 0:
            # Get information
            errors = pl_module.errors
            topology: ConnectionTopology = pl_module.topology_layer
-            components: Components = pl_module.proto_layer.components
+            components = pl_module.proto_layer.components

            # Insertion point
            worst = torch.argmax(errors)
@@ -121,8 +132,9 @@ class GNGCallback(pl.Callback):

            # Add component
            pl_module.proto_layer.add_components(
-                None,
-                initializer=LiteralCompInitializer(new_component.unsqueeze(0)))
+                1,
+                initializer=LiteralCompInitializer(new_component.unsqueeze(0)),
+            )

            # Adjust Topology
            topology.add_prototype()
--- a/prototorch/models/cbc.py
+++ b/prototorch/models/cbc.py
@@ -1,12 +1,12 @@
 import torch
 import torchmetrics
+from prototorch.core.competitions import CBCC
+from prototorch.core.components import ReasoningComponents
+from prototorch.core.initializers import RandomReasoningsInitializer
+from prototorch.core.losses import MarginLoss
+from prototorch.core.similarities import euclidean_similarity
+from prototorch.nn.wrappers import LambdaLayer

-from ..core.competitions import CBCC
-from ..core.components import ReasoningComponents
-from ..core.initializers import RandomReasoningsInitializer
-from ..core.losses import MarginLoss
-from ..core.similarities import euclidean_similarity
-from ..nn.wrappers import LambdaLayer
 from .abstract import ImagePrototypesMixin
 from .glvq import SiameseGLVQ

--- a/prototorch/models/extras.py
+++ b/prototorch/models/extras.py
@@ -5,8 +5,7 @@ Modules not yet available in prototorch go here temporarily.
 """

 import torch
-
-from ..core.similarities import gaussian
+from prototorch.core.similarities import gaussian


 def rank_scaled_gaussian(distances, lambd):
--- a/prototorch/models/glvq.py
+++ b/prototorch/models/glvq.py
@@ -1,22 +1,22 @@
 """Models based on the GLVQ framework."""

 import torch
-from torch.nn.parameter import Parameter
-
-from ..core.competitions import wtac
-from ..core.distances import (
+from prototorch.core.competitions import wtac
+from prototorch.core.distances import (
    lomega_distance,
    omega_distance,
    squared_euclidean_distance,
 )
-from ..core.initializers import EyeLinearTransformInitializer
-from ..core.losses import (
+from prototorch.core.initializers import EyeLinearTransformInitializer
+from prototorch.core.losses import (
    GLVQLoss,
    lvq1_loss,
    lvq21_loss,
 )
-from ..core.transforms import LinearTransform
-from ..nn.wrappers import LambdaLayer, LossLayer
+from prototorch.core.transforms import LinearTransform
+from prototorch.nn.wrappers import LambdaLayer, LossLayer
+from torch.nn.parameter import Parameter
+
 from .abstract import ImagePrototypesMixin, SupervisedPrototypeModel
 from .extras import ltangent_distance, orthogonalization

@@ -34,9 +34,9 @@ class GLVQ(SupervisedPrototypeModel):

        # Loss
        self.loss = GLVQLoss(
-            margin=self.hparams.margin,
-            transfer_fn=self.hparams.transfer_fn,
-            beta=self.hparams.transfer_beta,
+            margin=self.hparams["margin"],
+            transfer_fn=self.hparams["transfer_fn"],
+            beta=self.hparams["transfer_beta"],
        )

    # def on_save_checkpoint(self, checkpoint):
@@ -48,7 +48,7 @@ class GLVQ(SupervisedPrototypeModel):
            "prototype_win_ratios",
            torch.zeros(self.num_prototypes, device=self.device))

-    def on_epoch_start(self):
+    def on_train_epoch_start(self):
        self.initialize_prototype_win_ratios()

    def log_prototype_win_ratios(self, distances):
@@ -125,11 +125,11 @@ class SiameseGLVQ(GLVQ):

    def configure_optimizers(self):
        proto_opt = self.optimizer(self.proto_layer.parameters(),
-                                   lr=self.hparams.proto_lr)
+                                   lr=self.hparams["proto_lr"])
        # Only add a backbone optimizer if backbone has trainable parameters
        bb_params = list(self.backbone.parameters())
        if (bb_params):
-            bb_opt = self.optimizer(bb_params, lr=self.hparams.bb_lr)
+            bb_opt = self.optimizer(bb_params, lr=self.hparams["bb_lr"])
            optimizers = [proto_opt, bb_opt]
        else:
            optimizers = [proto_opt]
@@ -199,12 +199,13 @@ class GRLVQ(SiameseGLVQ):
    TODO Make a RelevanceLayer. `bb_lr` is ignored otherwise.

    """
+    _relevances: torch.Tensor

    def __init__(self, hparams, **kwargs):
        super().__init__(hparams, **kwargs)

        # Additional parameters
-        relevances = torch.ones(self.hparams.input_dim, device=self.device)
+        relevances = torch.ones(self.hparams["input_dim"], device=self.device)
        self.register_parameter("_relevances", Parameter(relevances))

        # Override the backbone
@@ -233,8 +234,8 @@ class SiameseGMLVQ(SiameseGLVQ):
        omega_initializer = kwargs.get("omega_initializer",
                                       EyeLinearTransformInitializer())
        self.backbone = LinearTransform(
-            self.hparams.input_dim,
-            self.hparams.latent_dim,
+            self.hparams["input_dim"],
+            self.hparams["latent_dim"],
            initializer=omega_initializer,
        )

@@ -244,7 +245,7 @@ class SiameseGMLVQ(SiameseGLVQ):

    @property
    def lambda_matrix(self):
-        omega = self.backbone.weight  # (input_dim, latent_dim)
+        omega = self.backbone.weights  # (input_dim, latent_dim)
        lam = omega @ omega.T
        return lam.detach().cpu()

@@ -257,6 +258,9 @@ class GMLVQ(GLVQ):

    """

+    # Parameters
+    _omega: torch.Tensor
+
    def __init__(self, hparams, **kwargs):
        distance_fn = kwargs.pop("distance_fn", omega_distance)
        super().__init__(hparams, distance_fn=distance_fn, **kwargs)
@@ -264,8 +268,8 @@ class GMLVQ(GLVQ):
        # Additional parameters
        omega_initializer = kwargs.get("omega_initializer",
                                       EyeLinearTransformInitializer())
-        omega = omega_initializer.generate(self.hparams.input_dim,
-                                           self.hparams.latent_dim)
+        omega = omega_initializer.generate(self.hparams["input_dim"],
+                                           self.hparams["latent_dim"])
        self.register_parameter("_omega", Parameter(omega))
        self.backbone = LambdaLayer(lambda x: x @ self._omega,
                                    name="omega matrix")
@@ -299,8 +303,8 @@ class LGMLVQ(GMLVQ):
        # Re-register `_omega` to override the one from the super class.
        omega = torch.randn(
            self.num_prototypes,
-            self.hparams.input_dim,
-            self.hparams.latent_dim,
+            self.hparams["input_dim"],
+            self.hparams["latent_dim"],
            device=self.device,
        )
        self.register_parameter("_omega", Parameter(omega))
@@ -316,23 +320,27 @@ class GTLVQ(LGMLVQ):
        omega_initializer = kwargs.get("omega_initializer")

        if omega_initializer is not None:
-            subspace = omega_initializer.generate(self.hparams.input_dim,
-                                                  self.hparams.latent_dim)
-            omega = torch.repeat_interleave(subspace.unsqueeze(0),
-                                            self.num_prototypes,
-                                            dim=0)
+            subspace = omega_initializer.generate(
+                self.hparams["input_dim"],
+                self.hparams["latent_dim"],
+            )
+            omega = torch.repeat_interleave(
+                subspace.unsqueeze(0),
+                self.num_prototypes,
+                dim=0,
+            )
        else:
            omega = torch.rand(
                self.num_prototypes,
-                self.hparams.input_dim,
-                self.hparams.latent_dim,
+                self.hparams["input_dim"],
+                self.hparams["latent_dim"],
                device=self.device,
            )

        # Re-register `_omega` to override the one from the super class.
        self.register_parameter("_omega", Parameter(omega))

-    def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
+    def on_train_batch_end(self, outputs, batch, batch_idx):
        with torch.no_grad():
            self._omega.copy_(orthogonalization(self._omega))

@@ -389,7 +397,7 @@ class ImageGTLVQ(ImagePrototypesMixin, GTLVQ):

    """

-    def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
+    def on_train_batch_end(self, outputs, batch, batch_idx):
        """Constrain the components to the range [0, 1] by clamping after updates."""
        self.proto_layer.components.data.clamp_(0.0, 1.0)
        with torch.no_grad():
--- a/prototorch/models/knn.py
+++ b/prototorch/models/knn.py
@@ -2,13 +2,14 @@

 import warnings

-from ..core.competitions import KNNC
-from ..core.components import LabeledComponents
-from ..core.initializers import (
+from prototorch.core.competitions import KNNC
+from prototorch.core.components import LabeledComponents
+from prototorch.core.initializers import (
    LiteralCompInitializer,
    LiteralLabelsInitializer,
 )
-from ..utils.utils import parse_data_arg
+from prototorch.utils.utils import parse_data_arg
+
 from .abstract import SupervisedPrototypeModel


@@ -36,10 +37,7 @@ class KNN(SupervisedPrototypeModel):
    def training_step(self, train_batch, batch_idx, optimizer_idx=None):
        return 1  # skip training step

-    def on_train_batch_start(self,
-                             train_batch,
-                             batch_idx,
-                             dataloader_idx=None):
+    def on_train_batch_start(self, train_batch, batch_idx):
        warnings.warn("k-NN has no training, skipping!")
        return -1

--- a/prototorch/models/lvq.py
+++ b/prototorch/models/lvq.py
@@ -1,8 +1,11 @@
 """LVQ models that are optimized using non-gradient methods."""

-from ..core.losses import _get_dp_dm
-from ..nn.activations import get_activation
-from ..nn.wrappers import LambdaLayer
+import logging
+
+from prototorch.core.losses import _get_dp_dm
+from prototorch.nn.activations import get_activation
+from prototorch.nn.wrappers import LambdaLayer
+
 from .abstract import NonGradientMixin
 from .glvq import GLVQ

@@ -29,8 +32,8 @@ class LVQ1(NonGradientMixin, GLVQ):
            self.proto_layer.load_state_dict({"_components": updated_protos},
                                             strict=False)

-        print(f"dis={dis}")
-        print(f"y={y}")
+        logging.debug(f"dis={dis}")
+        logging.debug(f"y={y}")
        # Logging
        self.log_acc(dis, y, tag="train_acc")

@@ -73,8 +76,7 @@ class MedianLVQ(NonGradientMixin, GLVQ):

    """

-    def __init__(self, hparams, verbose=True, **kwargs):
-        self.verbose = verbose
+    def __init__(self, hparams, **kwargs):
        super().__init__(hparams, **kwargs)

        self.transfer_layer = LambdaLayer(
@@ -115,8 +117,7 @@ class MedianLVQ(NonGradientMixin, GLVQ):
                _protos[i] = xk
                _lower_bound = self.lower_bound(x, y, _protos, plabels, gamma)
                if _lower_bound > lower_bound:
-                    if self.verbose:
-                        print(f"Updating prototype {i} to data {k}...")
+                    logging.debug(f"Updating prototype {i} to data {k}...")
                    self.proto_layer.load_state_dict({"_components": _protos},
                                                     strict=False)
                    break
--- a/prototorch/models/probabilistic.py
+++ b/prototorch/models/probabilistic.py
@@ -1,10 +1,13 @@
 """Probabilistic GLVQ methods"""

 import torch
+from prototorch.core.losses import nllr_loss, rslvq_loss
+from prototorch.core.pooling import (
+    stratified_min_pooling,
+    stratified_sum_pooling,
+)
+from prototorch.nn.wrappers import LossLayer

-from ..core.losses import nllr_loss, rslvq_loss
-from ..core.pooling import stratified_min_pooling, stratified_sum_pooling
-from ..nn.wrappers import LambdaLayer, LossLayer
 from .extras import GaussianPrior, RankScaledGaussianPrior
 from .glvq import GLVQ, SiameseGMLVQ

@@ -34,17 +37,24 @@ class ProbabilisticLVQ(GLVQ):
    def __init__(self, hparams, rejection_confidence=0.0, **kwargs):
        super().__init__(hparams, **kwargs)

-        self.conditional_distribution = None
        self.rejection_confidence = rejection_confidence
+        self._conditional_distribution = None

    def forward(self, x):
        distances = self.compute_distances(x)
+
        conditional = self.conditional_distribution(distances)
        prior = (1. / self.num_prototypes) * torch.ones(self.num_prototypes,
                                                        device=self.device)
        posterior = conditional * prior
+
        plabels = self.proto_layer._labels
-        y_pred = stratified_sum_pooling(posterior, plabels)
+        if isinstance(plabels, torch.LongTensor) or isinstance(
+                plabels, torch.cuda.LongTensor):  # type: ignore
+            y_pred = stratified_sum_pooling(posterior, plabels)  # type: ignore
+        else:
+            raise ValueError("Labels must be LongTensor.")
+
        return y_pred

    def predict(self, x):
@@ -61,6 +71,12 @@ class ProbabilisticLVQ(GLVQ):
        loss = batch_loss.sum()
        return loss

+    def conditional_distribution(self, distances):
+        """Conditional distribution of distances."""
+        if self._conditional_distribution is None:
+            raise ValueError("Conditional distribution is not set.")
+        return self._conditional_distribution(distances)
+

 class SLVQ(ProbabilisticLVQ):
    """Soft Learning Vector Quantization."""
@@ -72,7 +88,7 @@ class SLVQ(ProbabilisticLVQ):
        self.hparams.setdefault("variance", 1.0)
        variance = self.hparams.get("variance")

-        self.conditional_distribution = GaussianPrior(variance)
+        self._conditional_distribution = GaussianPrior(variance)
        self.loss = LossLayer(nllr_loss)


@@ -86,7 +102,7 @@ class RSLVQ(ProbabilisticLVQ):
        self.hparams.setdefault("variance", 1.0)
        variance = self.hparams.get("variance")

-        self.conditional_distribution = GaussianPrior(variance)
+        self._conditional_distribution = GaussianPrior(variance)
        self.loss = LossLayer(rslvq_loss)


--- a/prototorch/models/unsupervised.py
+++ b/prototorch/models/unsupervised.py
@@ -2,11 +2,10 @@

 import numpy as np
 import torch
+from prototorch.core.competitions import wtac
+from prototorch.core.distances import squared_euclidean_distance
+from prototorch.core.losses import NeuralGasEnergy

-from ..core.competitions import wtac
-from ..core.distances import squared_euclidean_distance
-from ..core.losses import NeuralGasEnergy
-from ..nn.wrappers import LambdaLayer
 from .abstract import NonGradientMixin, UnsupervisedPrototypeModel
 from .callbacks import GNGCallback
 from .extras import ConnectionTopology
@@ -18,6 +17,7 @@ class KohonenSOM(NonGradientMixin, UnsupervisedPrototypeModel):
    TODO Allow non-2D grids

    """
+    _grid: torch.Tensor

    def __init__(self, hparams, **kwargs):
        h, w = hparams.get("shape")
@@ -93,10 +93,10 @@ class NeuralGas(UnsupervisedPrototypeModel):
        self.hparams.setdefault("age_limit", 10)
        self.hparams.setdefault("lm", 1)

-        self.energy_layer = NeuralGasEnergy(lm=self.hparams.lm)
+        self.energy_layer = NeuralGasEnergy(lm=self.hparams["lm"])
        self.topology_layer = ConnectionTopology(
-            agelimit=self.hparams.age_limit,
-            num_prototypes=self.hparams.num_prototypes,
+            agelimit=self.hparams["age_limit"],
+            num_prototypes=self.hparams["num_prototypes"],
        )

    def training_step(self, train_batch, batch_idx):
@@ -109,12 +109,9 @@ class NeuralGas(UnsupervisedPrototypeModel):
        self.log("loss", loss)
        return loss

-    # def training_epoch_end(self, training_step_outputs):
-    #     print(f"{self.trainer.lr_schedulers}")
-    #     print(f"{self.trainer.lr_schedulers[0]['scheduler'].optimizer}")
-

 class GrowingNeuralGas(NeuralGas):
+    errors: torch.Tensor

    def __init__(self, hparams, **kwargs):
        super().__init__(hparams, **kwargs)
@@ -124,7 +121,10 @@ class GrowingNeuralGas(NeuralGas):
        self.hparams.setdefault("insert_reduction", 0.1)
        self.hparams.setdefault("insert_freq", 10)

-        errors = torch.zeros(self.hparams.num_prototypes, device=self.device)
+        errors = torch.zeros(
+            self.hparams["num_prototypes"],
+            device=self.device,
+        )
        self.register_buffer("errors", errors)

    def training_step(self, train_batch, _batch_idx):
@@ -139,7 +139,7 @@ class GrowingNeuralGas(NeuralGas):
        dp = d * mask

        self.errors += torch.sum(dp * dp)
-        self.errors *= self.hparams.step_reduction
+        self.errors *= self.hparams["step_reduction"]

        self.topology_layer(d)
        self.log("loss", loss)
@@ -148,7 +148,7 @@ class GrowingNeuralGas(NeuralGas):
    def configure_callbacks(self):
        return [
            GNGCallback(
-                reduction=self.hparams.insert_reduction,
-                freq=self.hparams.insert_freq,
+                reduction=self.hparams["insert_reduction"],
+                freq=self.hparams["insert_freq"],
            )
        ]
--- a/prototorch/models/vis.py
+++ b/prototorch/models/vis.py
@@ -1,15 +1,18 @@
 """Visualization Callbacks."""

+import warnings
+from typing import Sized
+
 import numpy as np
 import pytorch_lightning as pl
 import torch
 import torchvision
 from matplotlib import pyplot as plt
+from prototorch.utils.colors import get_colors, get_legend_handles
+from prototorch.utils.utils import mesh2d
+from pytorch_lightning.loggers import TensorBoardLogger
 from torch.utils.data import DataLoader, Dataset

-from ..utils.colors import get_colors, get_legend_handles
-from ..utils.utils import mesh2d
-

 class Vis2DAbstract(pl.Callback):

@@ -34,8 +37,13 @@ class Vis2DAbstract(pl.Callback):

        if data:
            if isinstance(data, Dataset):
-                x, y = next(iter(DataLoader(data, batch_size=len(data))))
-            elif isinstance(data, torch.utils.data.DataLoader):
+                if isinstance(data, Sized):
+                    x, y = next(iter(DataLoader(data, batch_size=len(data))))
+                else:
+                    # TODO: Add support for non-sized datasets
+                    raise NotImplementedError(
+                        "Data must be a dataset with a __len__ method.")
+            elif isinstance(data, DataLoader):
                x = torch.tensor([])
                y = torch.tensor([])
                for x_b, y_b in data:
@@ -123,7 +131,7 @@ class Vis2DAbstract(pl.Callback):
            else:
                plt.show(block=self.block)

-    def on_epoch_end(self, trainer, pl_module):
+    def on_train_epoch_end(self, trainer, pl_module):
        if not self.precheck(trainer):
            return True
        self.visualize(pl_module)
@@ -132,6 +140,9 @@ class Vis2DAbstract(pl.Callback):
    def on_train_end(self, trainer, pl_module):
        plt.close()

+    def visualize(self, pl_module):
+        raise NotImplementedError
+

 class VisGLVQ2D(Vis2DAbstract):

@@ -292,30 +303,45 @@ class VisImgComp(Vis2DAbstract):
        self.add_embedding = add_embedding
        self.embedding_data = embedding_data

-    def on_train_start(self, trainer, pl_module):
-        tb = pl_module.logger.experiment
-        if self.add_embedding:
-            ind = np.random.choice(len(self.x_train),
-                                   size=self.embedding_data,
-                                   replace=False)
-            data = self.x_train[ind]
-            tb.add_embedding(data.view(len(ind), -1),
-                             label_img=data,
-                             global_step=None,
-                             tag="Data Embedding",
-                             metadata=self.y_train[ind],
-                             metadata_header=None)
+    def on_train_start(self, _, pl_module):
+        if isinstance(pl_module.logger, TensorBoardLogger):
+            tb = pl_module.logger.experiment

-        if self.random_data:
-            ind = np.random.choice(len(self.x_train),
-                                   size=self.random_data,
-                                   replace=False)
-            data = self.x_train[ind]
-            grid = torchvision.utils.make_grid(data, nrow=self.num_columns)
-            tb.add_image(tag="Data",
-                         img_tensor=grid,
-                         global_step=None,
-                         dataformats=self.dataformats)
+            # Add embedding
+            if self.add_embedding:
+                if self.x_train is not None and self.y_train is not None:
+                    ind = np.random.choice(len(self.x_train),
+                                           size=self.embedding_data,
+                                           replace=False)
+                    data = self.x_train[ind]
+                    tb.add_embedding(data.view(len(ind), -1),
+                                     label_img=data,
+                                     global_step=None,
+                                     tag="Data Embedding",
+                                     metadata=self.y_train[ind],
+                                     metadata_header=None)
+                else:
+                    raise ValueError("No data for add embedding flag")
+
+            # Random Data
+            if self.random_data:
+                if self.x_train is not None:
+                    ind = np.random.choice(len(self.x_train),
+                                           size=self.random_data,
+                                           replace=False)
+                    data = self.x_train[ind]
+                    grid = torchvision.utils.make_grid(data,
+                                                       nrow=self.num_columns)
+                    tb.add_image(tag="Data",
+                                 img_tensor=grid,
+                                 global_step=None,
+                                 dataformats=self.dataformats)
+                else:
+                    raise ValueError("No data for random data flag")
+
+        else:
+            warnings.warn(
+                f"TensorBoardLogger is required, got {type(pl_module.logger)}")

    def add_to_tensorboard(self, trainer, pl_module):
        tb = pl_module.logger.experiment
--- a/setup.py
+++ b/setup.py
@@ -25,6 +25,7 @@ INSTALL_REQUIRES = [
    "prototorch>=0.7.3",
    "pytorch_lightning>=1.6.0",
    "torchmetrics",
+    "protobuf<3.20.0",
 ]
 CLI = [
    "jsonargparse",
@@ -54,7 +55,7 @@ ALL = CLI + DEV + DOCS + EXAMPLES + TESTS

 setup(
    name=safe_name("prototorch_" + PLUGIN_NAME),
-    version="0.5.0",
+    version="0.5.2",
    description="Pre-packaged prototype-based "
    "machine learning models using ProtoTorch and PyTorch-Lightning.",
    long_description=long_description,
@@ -80,6 +81,7 @@ setup(
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: MIT License",
        "Natural Language :: English",
+        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.8",
Author	SHA1	Message	Date
Alexander Engelsberger	fcd944d3ff	build: bump version 0.5.1 → 0.5.2	2022-06-01 14:25:44 +02:00
Alexander Engelsberger	054720dd7b	fix(hotfix): Protobuf error workaround	2022-06-01 14:14:57 +02:00
Alexander Engelsberger	d16a0de202	build: bump version 0.5.0 → 0.5.1	2022-05-17 12:04:08 +02:00
Alexander Engelsberger	76fea3f881	chore: update all examples to pytorch 1.6	2022-05-17 12:03:43 +02:00
Alexander Engelsberger	c00513ae0d	chore: minor updates and version updates	2022-05-17 12:00:52 +02:00
Alexander Engelsberger	bccef8bef0	chore: replace relative imports	2022-05-16 11:12:53 +02:00
Alexander Engelsberger	29ee326b85	ci: Update PreCommit hooks	2022-05-16 11:11:48 +02:00
Jensun Ravichandran	055568dc86	fix: glvq_iris example works again	2022-05-09 17:33:52 +02:00
Alexander Engelsberger	3a7328e290	chore: small changes	2022-04-27 10:37:12 +02:00