build: bump version 0.7.0 → 0.7.1

.bumpversion.cfg

@@ -1,13 +1,13 @@
 [bumpversion]
-current_version = 0.5.0
+current_version = 0.7.1
 commit = True
 tag = True
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)
 serialize = {major}.{minor}.{patch}
 message = build: bump version {current_version} → {new_version}
 
-[bumpversion:file:setup.py]
+[bumpversion:file:pyproject.toml]
 
-[bumpversion:file:./prototorch/models/__init__.py]
+[bumpversion:file:./src/prototorch/models/__init__.py]
 
 [bumpversion:file:./docs/source/conf.py]

.github/workflows/examples.yml

@@ -6,20 +6,20 @@ name: examples
 on:
   push:
     paths:
-      - 'examples/**.py'
+      - "examples/**.py"
 jobs:
   cpu:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
-    - name: Set up Python 3.10
+      - name: Set up Python 3.11
-      uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
-      with:
+        with:
-        python-version: "3.10"
+          python-version: "3.11"
-    - name: Install dependencies
+      - name: Install dependencies
-      run: |
+        run: |
-        python -m pip install --upgrade pip
+          python -m pip install --upgrade pip
-        pip install .[all]
+          pip install .[all]
-    - name: Run examples
+      - name: Run examples
-      run: |
+        run: |
-        ./tests/test_examples.sh examples/
+          ./tests/test_examples.sh examples/

.github/workflows/pythonapp.yml

@@ -6,70 +6,70 @@ name: tests
 on:
   push:
   pull_request:
-    branches: [ master ]
+    branches: [master]
 
 jobs:
   style:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
-    - name: Set up Python 3.10
+      - name: Set up Python 3.11
-      uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
-      with:
+        with:
-        python-version: "3.10"
+          python-version: "3.11"
-    - name: Install dependencies
+      - name: Install dependencies
-      run: |
+        run: |
-        python -m pip install --upgrade pip
+          python -m pip install --upgrade pip
-        pip install .[all]
+          pip install .[all]
-    - uses: pre-commit/action@v2.0.3
+      - uses: pre-commit/action@v3.0.0
   compatibility:
     needs: style
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         os: [ubuntu-latest, windows-latest]
         exclude:
-        - os: windows-latest
+          - os: windows-latest
-          python-version: "3.7"
+            python-version: "3.8"
-        - os: windows-latest
+          - os: windows-latest
-          python-version: "3.8"
+            python-version: "3.9"
-        - os: windows-latest
+          - os: windows-latest
-          python-version: "3.9"
+            python-version: "3.10"
 
     runs-on: ${{ matrix.os }}
     steps:
-    - uses: actions/checkout@v2
+      - uses: actions/checkout@v2
-    - name: Set up Python ${{ matrix.python-version }}
+      - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
-      with:
+        with:
-        python-version: ${{ matrix.python-version }}
+          python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
+      - name: Install dependencies
-      run: |
+        run: |
-        python -m pip install --upgrade pip
+          python -m pip install --upgrade pip
-        pip install .[all]
+          pip install .[all]
-    - name: Test with pytest
+      - name: Test with pytest
-      run: |
+        run: |
-        pytest
+          pytest
   publish_pypi:
     if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
     needs: compatibility
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
-    - name: Set up Python 3.10
+      - name: Set up Python 3.11
-      uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
-      with:
+        with:
-        python-version: "3.10"
+          python-version: "3.11"
-    - name: Install dependencies
+      - name: Install dependencies
-      run: |
+        run: |
-        python -m pip install --upgrade pip
+          python -m pip install --upgrade pip
-        pip install .[all]
+          pip install .[all]
-        pip install wheel
+          pip install build
-    - name: Build package
+      - name: Build package
-      run: python setup.py sdist bdist_wheel
+        run: python -m build . -C verbose
-    - name: Publish a Python distribution to PyPI
+      - name: Publish a Python distribution to PyPI
-      uses: pypa/gh-action-pypi-publish@release/v1
+        uses: pypa/gh-action-pypi-publish@release/v1
-      with:
+        with:
-        user: __token__
+          user: __token__
-        password: ${{ secrets.PYPI_API_TOKEN }}
+          password: ${{ secrets.PYPI_API_TOKEN }}

.pre-commit-config.yaml

@@ -2,52 +2,53 @@
 # See https://pre-commit.com/hooks.html for more hooks
 
 repos:
-- repo: https://github.com/pre-commit/pre-commit-hooks
+  - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.1.0
+    rev: v4.4.0
-  hooks:
+    hooks:
-  - id: trailing-whitespace
+      - id: trailing-whitespace
-  - id: end-of-file-fixer
+      - id: end-of-file-fixer
-  - id: check-yaml
+      - id: check-yaml
-  - id: check-added-large-files
+      - id: check-added-large-files
-  - id: check-ast
+      - id: check-ast
-  - id: check-case-conflict
+      - id: check-case-conflict
 
-- repo: https://github.com/myint/autoflake
+  - repo: https://github.com/myint/autoflake
-  rev: v1.4
+    rev: v2.1.1
-  hooks:
+    hooks:
-  - id: autoflake
+      - id: autoflake
 
-- repo: http://github.com/PyCQA/isort
+  - repo: http://github.com/PyCQA/isort
-  rev: 5.10.1
+    rev: 5.12.0
-  hooks:
+    hooks:
-  - id: isort
+      - id: isort
 
-- repo: https://github.com/pre-commit/mirrors-mypy
+  - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.931
+    rev: v1.3.0
-  hooks:
+    hooks:
-  - id: mypy
+      - id: mypy
-    files: prototorch
+        files: prototorch
-    additional_dependencies: [types-pkg_resources]
+        additional_dependencies: [types-pkg_resources]
 
-- repo: https://github.com/pre-commit/mirrors-yapf
+  - repo: https://github.com/pre-commit/mirrors-yapf
-  rev: v0.32.0
+    rev: v0.32.0
-  hooks:
+    hooks:
-  - id: yapf
+      - id: yapf
+        additional_dependencies: ["toml"]
 
-- repo: https://github.com/pre-commit/pygrep-hooks
+  - repo: https://github.com/pre-commit/pygrep-hooks
-  rev: v1.9.0
+    rev: v1.10.0
-  hooks:
+    hooks:
-  - id: python-use-type-annotations
+      - id: python-use-type-annotations
-  - id: python-no-log-warn
+      - id: python-no-log-warn
-  - id: python-check-blanket-noqa
+      - id: python-check-blanket-noqa
 
-- repo: https://github.com/asottile/pyupgrade
+  - repo: https://github.com/asottile/pyupgrade
-  rev: v2.31.0
+    rev: v3.7.0
-  hooks:
+    hooks:
-  - id: pyupgrade
+      - id: pyupgrade
 
-- repo: https://github.com/si-cim/gitlint
+  - repo: https://github.com/si-cim/gitlint
-  rev: v0.15.2-unofficial
+    rev: v0.15.2-unofficial
-  hooks:
+    hooks:
-  - id: gitlint
+      - id: gitlint
-    args: [--contrib=CT1, --ignore=B6, --msg-filename]
+        args: [--contrib=CT1, --ignore=B6, --msg-filename]

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.7.1"
 
 # -- General configuration ---------------------------------------------------
 

examples/cbc_iris.py

@@ -1,25 +1,32 @@
 """CBC example using the Iris dataset."""
-
 import argparse
+import warnings
 
 import prototorch as pt
 import pytorch_lightning as pl
-import torch
+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)
 
 if __name__ == "__main__":
+    # Reproducibility
+    seed_everything(seed=4)
+
     # Command-line arguments
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # 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,32 @@ if __name__ == "__main__":
     )
 
     # Initialize the model
-    model = pt.models.CBC(
+    model = CBC(
         hparams,
-        components_initializer=pt.initializers.SSCI(train_ds, noise=0.01),
+        components_initializer=pt.initializers.SSCI(train_ds, noise=0.1),
-        reasonings_iniitializer=pt.initializers.
+        reasonings_initializer=pt.initializers.
         PurePositiveReasoningsInitializer(),
     )
 
     # Callbacks
-    vis = pt.models.VisCBC2D(data=train_ds,
+    vis = VisCBC2D(
-                             title="CBC Iris Example",
+        data=train_ds,
-                             resolution=100,
+        title="CBC Iris Example",
-                             axis_off=True)
+        resolution=100,
+        axis_off=True,
+    )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
-        callbacks=[vis],
+        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,
     )
 
     # Training loop

examples/dynamic_pruning.py

@@ -1,30 +1,50 @@
 """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 lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import (
+    CELVQ,
+    PruneLoserPrototypes,
+    VisGLVQ2D,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+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)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Dataset
     num_classes = 4
     num_features = 2
     num_clusters = 1
-    train_ds = pt.datasets.Random(num_samples=500,
+    train_ds = pt.datasets.Random(
-                                  num_classes=num_classes,
+        num_samples=500,
-                                  num_features=num_features,
+        num_classes=num_classes,
-                                  num_clusters=num_clusters,
+        num_features=num_features,
-                                  separation=3.0,
+        num_clusters=num_clusters,
-                                  seed=42)
+        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 +54,7 @@ if __name__ == "__main__":
     )
 
     # Initialize the model
-    model = pt.models.CELVQ(
+    model = CELVQ(
         hparams,
         prototypes_initializer=pt.initializers.FVCI(2, 3.0),
     )
@@ -43,18 +63,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)
+    vis = VisGLVQ2D(train_ds)
-    pruning = pt.models.PruneLoserPrototypes(
+    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,
@@ -64,17 +84,18 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         callbacks=[
             vis,
             pruning,
             es,
         ],
-        progress_bar_refresh_rate=0,
+        detect_anomaly=True,
-        terminate_on_nan=True,
+        log_every_n_steps=1,
-        weights_summary="full",
+        max_epochs=1000,
-        accelerator="ddp",
     )
 
     # Training loop

examples/glvq_iris.py

@@ -1,23 +1,35 @@
 """GLVQ example using the Iris dataset."""
 
 import argparse
+import logging
+import warnings
 
 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import GLVQ, VisGLVQ2D
+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)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # 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, num_workers=4)
 
     # Hyperparameters
     hparams = dict(
@@ -29,7 +41,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 +53,19 @@ 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(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
-        callbacks=[vis],
+        devices=args.gpus if args.gpus else "auto",
-        weights_summary="full",
+        fast_dev_run=args.fast_dev_run,
-        accelerator="ddp",
+        callbacks=[
+            vis,
+        ],
+        max_epochs=100,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop
@@ -58,8 +75,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)

examples/gmlvq_iris.py

@@ -1,23 +1,36 @@
 """GMLVQ example using the Iris dataset."""
 
 import argparse
+import warnings
 
 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import GMLVQ, VisGMLVQ2D
+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)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Dataset
     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 +45,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,15 +57,22 @@ 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(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
-        callbacks=[vis],
+        devices=args.gpus if args.gpus else "auto",
-        weights_summary="full",
+        fast_dev_run=args.fast_dev_run,
-        accelerator="ddp",
+        callbacks=[
+            vis,
+        ],
+        max_epochs=100,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop
     trainer.fit(model, train_loader)
+
+    torch.save(model, "iris.pth")

examples/gmlvq_mnist.py

@@ -1,17 +1,33 @@
 """GMLVQ example using the MNIST dataset."""
 
 import argparse
+import warnings
 
 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import (
+    ImageGMLVQ,
+    PruneLoserPrototypes,
+    VisImgComp,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+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)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Dataset
@@ -33,12 +49,8 @@ if __name__ == "__main__":
     )
 
     # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds,
+    train_loader = DataLoader(train_ds, num_workers=4, batch_size=256)
-                                               num_workers=0,
+    test_loader = DataLoader(test_ds, num_workers=4, batch_size=256)
-                                               batch_size=256)
-    test_loader = torch.utils.data.DataLoader(test_ds,
-                                              num_workers=0,
-                                              batch_size=256)
 
     # Hyperparameters
     num_classes = 10
@@ -52,14 +64,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 +81,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,
@@ -85,16 +97,18 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         callbacks=[
             vis,
             pruning,
-            # es,
+            es,
         ],
-        terminate_on_nan=True,
+        max_epochs=1000,
-        weights_summary=None,
+        log_every_n_steps=1,
-        # accelerator="ddp",
+        detect_anomaly=True,
     )
 
     # Training loop

examples/gmlvq_spiral.py

@@ -1,22 +1,39 @@
 """GMLVQ example using the spiral dataset."""
 
 import argparse
+import warnings
 
 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import (
+    GMLVQ,
+    PruneLoserPrototypes,
+    VisGLVQ2D,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+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)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Dataset
     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 +49,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 +70,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,
@@ -62,14 +79,18 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         callbacks=[
             vis,
             es,
             pruning,
         ],
-        terminate_on_nan=True,
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

examples/gng_iris.py

@@ -1,23 +1,33 @@
 """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 lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import GrowingNeuralGas, VisNG2D
+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
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     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 +37,7 @@ if __name__ == "__main__":
     )
 
     # Initialize the model
-    model = pt.models.GrowingNeuralGas(
+    model = GrowingNeuralGas(
         hparams,
         prototypes_initializer=pt.initializers.ZCI(2),
     )
@@ -36,17 +46,22 @@ 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(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
+        callbacks=[
+            vis,
+        ],
         max_epochs=100,
-        callbacks=[vis],
+        log_every_n_steps=1,
-        weights_summary="full",
+        detect_anomaly=True,
     )
 
     # Training loop

examples/grlvq_iris.py

@@ -0,0 +1,77 @@
+"""GMLVQ example using the Iris dataset."""
+
+import argparse
+import warnings
+
+import prototorch as pt
+import pytorch_lightning as pl
+import torch
+from lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import GRLVQ, VisSiameseGLVQ2D
+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.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
+    args = parser.parse_args()
+
+    # Dataset
+    train_ds = pt.datasets.Iris([0, 1])
+
+    # Dataloaders
+    train_loader = DataLoader(train_ds, batch_size=64)
+
+    # Hyperparameters
+    hparams = dict(
+        input_dim=2,
+        distribution={
+            "num_classes": 3,
+            "per_class": 2
+        },
+        proto_lr=0.01,
+        bb_lr=0.01,
+    )
+
+    # Initialize the model
+    model = GRLVQ(
+        hparams,
+        optimizer=torch.optim.Adam,
+        prototypes_initializer=pt.initializers.SMCI(train_ds),
+        lr_scheduler=ExponentialLR,
+        lr_scheduler_kwargs=dict(gamma=0.99, verbose=False),
+    )
+
+    # Compute intermediate input and output sizes
+    model.example_input_array = torch.zeros(4, 2)
+
+    # Callbacks
+    vis = VisSiameseGLVQ2D(data=train_ds)
+
+    # 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,
+        ],
+        max_epochs=5,
+        log_every_n_steps=1,
+        detect_anomaly=True,
+    )
+
+    # Training loop
+    trainer.fit(model, train_loader)
+
+    torch.save(model, "iris.pth")

examples/gtlvq_mnist.py

@@ -1,17 +1,34 @@
 """GTLVQ example using the MNIST dataset."""
 
 import argparse
+import warnings
 
 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import (
+    ImageGTLVQ,
+    PruneLoserPrototypes,
+    VisImgComp,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+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)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Dataset
@@ -33,12 +50,8 @@ if __name__ == "__main__":
     )
 
     # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds,
+    train_loader = DataLoader(train_ds, num_workers=0, batch_size=256)
-                                               num_workers=0,
+    test_loader = DataLoader(test_ds, num_workers=0, batch_size=256)
-                                               batch_size=256)
-    test_loader = torch.utils.data.DataLoader(test_ds,
-                                              num_workers=0,
-                                              batch_size=256)
 
     # Hyperparameters
     num_classes = 10
@@ -52,7 +65,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 +74,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 +84,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,
@@ -88,16 +101,18 @@ if __name__ == "__main__":
 
     # Setup trainer
     # using GPUs here is strongly recommended!
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         callbacks=[
             vis,
             pruning,
-            # es,
+            es,
         ],
-        terminate_on_nan=True,
+        max_epochs=1000,
-        weights_summary=None,
+        log_every_n_steps=1,
-        accelerator="ddp",
+        detect_anomaly=True,
     )
 
     # Training loop

examples/gtlvq_moons.py

@@ -1,45 +1,58 @@
 """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 lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import GTLVQ, VisGLVQ2D
+from pytorch_lightning.callbacks import EarlyStopping
+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
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     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,
+    train_loader = DataLoader(
-                                               batch_size=256,
+        train_ds,
-                                               shuffle=True)
+        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(
+    model = GTLVQ(hparams,
-        hparams, prototypes_initializer=pt.initializers.SMCI(train_ds))
+                  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)
+    vis = VisGLVQ2D(data=train_ds)
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
         monitor="train_acc",
         min_delta=0.001,
         patience=20,
@@ -49,14 +62,17 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         callbacks=[
             vis,
             es,
         ],
-        weights_summary="full",
+        max_epochs=1000,
-        accelerator="ddp",
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

examples/knn_iris.py

@@ -1,60 +1,75 @@
 """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
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # 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,
+    X_train, X_test, y_train, y_test = train_test_split(
-                                                        y,
+        X,
-                                                        test_size=0.5,
+        y,
-                                                        random_state=42)
+        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)
+    train_loader = DataLoader(train_ds, batch_size=16)
-    test_loader = torch.utils.data.DataLoader(test_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,
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         max_epochs=1,
-        callbacks=[vis],
+        callbacks=[
-        weights_summary="full",
+            vis,
+        ],
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop
@@ -63,7 +78,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)

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 lightning_fabric.utilities.seed import seed_everything
 from matplotlib import pyplot as plt
+from prototorch.models import KohonenSOM
 from prototorch.utils.colors import hex_to_rgb
+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],
+            plt.text(
-                     iloc[0],
+                iloc[1],
-                     cnames[i],
+                iloc[0],
-                     ha="center",
+                color_names[i],
-                     va="center",
+                ha="center",
-                     bbox=dict(facecolor="white", alpha=0.5, lw=0))
+                va="center",
+                bbox=dict(facecolor="white", alpha=0.5, lw=0),
+            )
 
         if trainer.current_epoch != trainer.max_epochs - 1:
             plt.pause(self.pause_time)
@@ -47,11 +58,12 @@ class Vis2DColorSOM(pl.Callback):
 if __name__ == "__main__":
     # Command-line arguments
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Reproducibility
-    pl.utilities.seed.seed_everything(seed=42)
+    seed_everything(seed=42)
 
     # Prepare the data
     hex_colors = [
@@ -59,15 +71,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_ds = TensorDataset(data)
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=8)
+    train_loader = DataLoader(train_ds, batch_size=8)
 
     # Hyperparameters
     hparams = dict(
@@ -78,7 +90,7 @@ if __name__ == "__main__":
     )
 
     # Initialize the model
-    model = pt.models.KohonenSOM(
+    model = KohonenSOM(
         hparams,
         prototypes_initializer=pt.initializers.RNCI(3),
     )
@@ -87,17 +99,22 @@ if __name__ == "__main__":
     model.example_input_array = torch.zeros(4, 3)
 
     # Model summary
-    print(model)
+    logging.info(model)
 
     # Callbacks
     vis = Vis2DColorSOM(data=data)
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         max_epochs=500,
-        callbacks=[vis],
+        callbacks=[
-        weights_summary="full",
+            vis,
+        ],
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

examples/lgmlvq_moons.py

@@ -1,27 +1,36 @@
 """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 lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import LGMLVQ, VisGLVQ2D
+from pytorch_lightning.callbacks import EarlyStopping
+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
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     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,
+    train_loader = DataLoader(train_ds, batch_size=256, shuffle=True)
-                                               batch_size=256,
-                                               shuffle=True)
 
     # Hyperparameters
     hparams = dict(
@@ -31,7 +40,7 @@ if __name__ == "__main__":
     )
 
     # Initialize the model
-    model = pt.models.LGMLVQ(
+    model = LGMLVQ(
         hparams,
         prototypes_initializer=pt.initializers.SMCI(train_ds),
     )
@@ -40,11 +49,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)
+    vis = VisGLVQ2D(data=train_ds)
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
         monitor="train_acc",
         min_delta=0.001,
         patience=20,
@@ -54,14 +63,17 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         callbacks=[
             vis,
             es,
         ],
-        weights_summary="full",
+        log_every_n_steps=1,
-        accelerator="ddp",
+        max_epochs=1000,
+        detect_anomaly=True,
     )
 
     # Training loop

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 lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import (
+    LVQMLN,
+    PruneLoserPrototypes,
+    VisSiameseGLVQ2D,
+)
+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):
@@ -27,17 +39,18 @@ class Backbone(torch.nn.Module):
 if __name__ == "__main__":
     # Command-line arguments
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Dataset
     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 +63,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 +72,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,
@@ -79,12 +89,17 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         callbacks=[
             vis,
             pruning,
         ],
+        log_every_n_steps=1,
+        max_epochs=1000,
+        detect_anomaly=True,
     )
 
     # Training loop

examples/median_lvq_iris.py

@@ -1,28 +1,40 @@
 """Median-LVQ example using the Iris dataset."""
 
 import argparse
+import warnings
 
 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import MedianLVQ, VisGLVQ2D
+from pytorch_lightning.callbacks import EarlyStopping
+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)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Dataset
     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 +43,8 @@ if __name__ == "__main__":
     model.example_input_array = torch.zeros(4, 2)
 
     # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
+    vis = VisGLVQ2D(data=train_ds)
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
         monitor="train_acc",
         min_delta=0.01,
         patience=5,
@@ -42,10 +54,17 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
-        callbacks=[vis, es],
+        devices=args.gpus if args.gpus else "auto",
-        weights_summary="full",
+        fast_dev_run=args.fast_dev_run,
+        callbacks=[
+            vis,
+            es,
+        ],
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

examples/ng_iris.py

@@ -1,23 +1,35 @@
 """Neural Gas example using the Iris dataset."""
 
 import argparse
+import warnings
 
 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import NeuralGas, VisNG2D
+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)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # 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 +37,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 +47,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 +57,20 @@ 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(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
-        callbacks=[vis],
+        devices=args.gpus if args.gpus else "auto",
-        weights_summary="full",
+        fast_dev_run=args.fast_dev_run,
+        callbacks=[
+            vis,
+        ],
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

examples/rslvq_iris.py

@@ -1,25 +1,34 @@
 """RSLVQ example using the Iris dataset."""
 
 import argparse
+import warnings
 
 import prototorch as pt
 import pytorch_lightning as pl
 import torch
+from lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import RSLVQ, VisGLVQ2D
+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
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     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 +42,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 +51,20 @@ 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(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
-        callbacks=[vis],
+        devices=args.gpus if args.gpus else "auto",
-        terminate_on_nan=True,
+        fast_dev_run=args.fast_dev_run,
-        weights_summary="full",
+        callbacks=[
-        accelerator="ddp",
+            vis,
+        ],
+        detect_anomaly=True,
+        max_epochs=100,
+        log_every_n_steps=1,
     )
 
     # Training loop

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 lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import SiameseGLVQ, VisSiameseGLVQ2D
+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):
@@ -27,46 +35,50 @@ class Backbone(torch.nn.Module):
 if __name__ == "__main__":
     # Command-line arguments
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Dataset
     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,
+        lr=0.01,
-        bb_lr=0.01,
     )
 
     # Initialize the backbone
     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(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
-        callbacks=[vis],
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
+        callbacks=[
+            vis,
+        ],
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

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 lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import SiameseGTLVQ, VisSiameseGLVQ2D
+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):
@@ -27,46 +35,52 @@ class Backbone(torch.nn.Module):
 if __name__ == "__main__":
     # Command-line arguments
     parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # Dataset
     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],
+    hparams = dict(
-                   proto_lr=0.01,
+        distribution=[1, 2, 3],
-                   bb_lr=0.01,
+        lr=0.01,
-                   input_dim=2,
+        input_dim=2,
-                   latent_dim=1)
+        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(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
-        callbacks=[vis],
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
+        callbacks=[
+            vis,
+        ],
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

examples/warm_starting.py

@@ -1,24 +1,42 @@
 """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 lightning_fabric.utilities.seed import seed_everything
+from prototorch.models import (
+    GLVQ,
+    KNN,
+    GrowingNeuralGas,
+    PruneLoserPrototypes,
+    VisGLVQ2D,
+)
+from pytorch_lightning.callbacks import EarlyStopping
+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)
+    parser.add_argument("--gpus", type=int, default=0)
+    parser.add_argument("--fast_dev_run", type=bool, default=False)
     args = parser.parse_args()
 
     # 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 +44,7 @@ if __name__ == "__main__":
     )
 
     # Callbacks
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
         monitor="loss",
         min_delta=0.001,
         patience=20,
@@ -37,9 +55,14 @@ if __name__ == "__main__":
 
     # Setup trainer for GNG
     trainer = pl.Trainer(
-        max_epochs=100,
+        accelerator="cpu",
-        callbacks=[es],
+        max_epochs=50 if args.fast_dev_run else
-        weights_summary=None,
+        1000,  # 10 epochs fast dev run reproducible DIV error.
+        callbacks=[
+            es,
+        ],
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop
@@ -52,12 +75,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 +93,15 @@ if __name__ == "__main__":
     model.example_input_array = torch.zeros(4, 2)
 
     # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
+    vis = VisGLVQ2D(data=train_ds)
-    pruning = pt.models.PruneLoserPrototypes(
+    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,
@@ -88,15 +111,18 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
+    trainer = pl.Trainer(
-        args,
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
         callbacks=[
             vis,
             pruning,
             es,
         ],
-        weights_summary="full",
+        max_epochs=1000,
-        accelerator="ddp",
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

pyproject.toml

@@ -0,0 +1,90 @@
+
+[project]
+name = "prototorch-models"
+version = "0.7.1"
+description = "Pre-packaged prototype-based machine learning models using ProtoTorch and PyTorch-Lightning."
+authors = [
+    { name = "Jensun Ravichandran", email = "jjensun@gmail.com" },
+    { name = "Alexander Engelsberger", email = "engelsbe@hs-mittweida.de" },
+]
+dependencies = ["lightning>=2.0.0", "prototorch>=0.7.5"]
+requires-python = ">=3.8"
+readme = "README.md"
+license = { text = "MIT" }
+classifiers = [
+    "Development Status :: 2 - Pre-Alpha",
+    "Environment :: Plugins",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Education",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Natural Language :: English",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+
+[project.urls]
+Homepage = "https://github.com/si-cim/prototorch_models"
+Downloads = "https://github.com/si-cim/prototorch_models.git"
+
+[project.optional-dependencies]
+dev = ["bumpversion", "pre-commit", "yapf", "toml"]
+examples = ["matplotlib", "scikit-learn"]
+ci = ["pytest", "pre-commit"]
+docs = [
+    "recommonmark",
+    "nbsphinx",
+    "sphinx",
+    "sphinx_rtd_theme",
+    "sphinxcontrib-bibtex",
+    "sphinxcontrib-katex",
+    "ipykernel",
+]
+all = [
+    "bumpversion",
+    "pre-commit",
+    "yapf",
+    "toml",
+    "pytest",
+    "matplotlib",
+    "scikit-learn",
+    "recommonmark",
+    "nbsphinx",
+    "sphinx",
+    "sphinx_rtd_theme",
+    "sphinxcontrib-bibtex",
+    "sphinxcontrib-katex",
+    "ipykernel",
+]
+
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.yapf]
+based_on_style = "pep8"
+spaces_before_comment = 2
+split_before_logical_operator = true
+
+[tool.pylint]
+disable = ["too-many-arguments", "too-few-public-methods", "fixme"]
+
+[tool.isort]
+profile = "hug"
+src_paths = ["isort", "test"]
+multi_line_output = 3
+include_trailing_comma = true
+force_grid_wrap = 3
+use_parentheses = true
+line_length = 79
+
+[tool.mypy]
+explicit_package_bases = true
+namespace_packages = true

setup.cfg

@@ -1,23 +0,0 @@
-[yapf]
-based_on_style = pep8
-spaces_before_comment = 2
-split_before_logical_operator = true
-
-[pylint]
-disable =
-	too-many-arguments,
-	too-few-public-methods,
-	fixme,
-
-
-[pycodestyle]
-max-line-length = 79
-
-[isort]
-profile = hug
-src_paths = isort, test
-multi_line_output = 3
-include_trailing_comma = True
-force_grid_wrap = 3
-use_parentheses = True
-line_length = 79

setup.py

@@ -1,97 +0,0 @@
-"""
-
- ######
- #     # #####   ####  #####  ####  #####  ####  #####   ####  #    #
- #     # #    # #    #   #   #    #   #   #    # #    # #    # #    #
- ######  #    # #    #   #   #    #   #   #    # #    # #      ######
- #       #####  #    #   #   #    #   #   #    # #####  #      #    #
- #       #   #  #    #   #   #    #   #   #    # #   #  #    # #    #
- #       #    #  ####    #    ####    #    ####  #    #  ####  #    #Plugin
-
-ProtoTorch models Plugin Package
-"""
-from pkg_resources import safe_name
-from setuptools import find_namespace_packages, setup
-
-PLUGIN_NAME = "models"
-
-PROJECT_URL = "https://github.com/si-cim/prototorch_models"
-DOWNLOAD_URL = "https://github.com/si-cim/prototorch_models.git"
-
-with open("README.md", "r") as fh:
-    long_description = fh.read()
-
-INSTALL_REQUIRES = [
-    "prototorch>=0.7.3",
-    "pytorch_lightning>=1.6.0",
-    "torchmetrics",
-]
-CLI = [
-    "jsonargparse",
-]
-DEV = [
-    "bumpversion",
-    "pre-commit",
-]
-DOCS = [
-    "recommonmark",
-    "sphinx",
-    "nbsphinx",
-    "ipykernel",
-    "sphinx_rtd_theme",
-    "sphinxcontrib-katex",
-    "sphinxcontrib-bibtex",
-]
-EXAMPLES = [
-    "matplotlib",
-    "scikit-learn",
-]
-TESTS = [
-    "codecov",
-    "pytest",
-]
-ALL = CLI + DEV + DOCS + EXAMPLES + TESTS
-
-setup(
-    name=safe_name("prototorch_" + PLUGIN_NAME),
-    version="0.5.0",
-    description="Pre-packaged prototype-based "
-    "machine learning models using ProtoTorch and PyTorch-Lightning.",
-    long_description=long_description,
-    long_description_content_type="text/markdown",
-    author="Alexander Engelsberger",
-    author_email="engelsbe@hs-mittweida.de",
-    url=PROJECT_URL,
-    download_url=DOWNLOAD_URL,
-    license="MIT",
-    python_requires=">=3.7",
-    install_requires=INSTALL_REQUIRES,
-    extras_require={
-        "dev": DEV,
-        "examples": EXAMPLES,
-        "tests": TESTS,
-        "all": ALL,
-    },
-    classifiers=[
-        "Development Status :: 2 - Pre-Alpha",
-        "Environment :: Plugins",
-        "Intended Audience :: Developers",
-        "Intended Audience :: Education",
-        "Intended Audience :: Science/Research",
-        "License :: OSI Approved :: MIT License",
-        "Natural Language :: English",
-        "Programming Language :: Python :: 3.10",
-        "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.7",
-        "Operating System :: OS Independent",
-        "Topic :: Scientific/Engineering :: Artificial Intelligence",
-        "Topic :: Software Development :: Libraries",
-        "Topic :: Software Development :: Libraries :: Python Modules",
-    ],
-    entry_points={
-        "prototorch.plugins": f"{PLUGIN_NAME} = prototorch.{PLUGIN_NAME}"
-    },
-    packages=find_namespace_packages(include=["prototorch.*"]),
-    zip_safe=False,
-)

src/prototorch/models/__init__.py

@@ -36,4 +36,4 @@ from .unsupervised import (
 )
 from .vis import *
 
-__version__ = "0.5.0"
+__version__ = "0.7.1"

src/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 prototorch.core.competitions import WTAC
-from ..core.competitions import WTAC
+from prototorch.core.components import (
-from ..core.components import Components, LabeledComponents
+    AbstractComponents,
-from ..core.distances import euclidean_distance
+    Components,
-from ..core.initializers import LabelsInitializer, ZerosCompInitializer
+    LabeledComponents,
-from ..core.pooling import stratified_min_pooling
+)
-from ..nn.wrappers import LambdaLayer
+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,12 +65,13 @@ class ProtoTorchBolt(pl.LightningModule):
 
 
 class PrototypeModel(ProtoTorchBolt):
+    proto_layer: AbstractComponents
 
     def __init__(self, hparams, **kwargs):
         super().__init__(hparams, **kwargs)
 
         distance_fn = kwargs.get("distance_fn", euclidean_distance)
-        self.distance_layer = LambdaLayer(distance_fn)
+        self.distance_layer = LambdaLayer(distance_fn, name="distance_fn")
 
     @property
     def num_prototypes(self):
@@ -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):
@@ -171,26 +186,37 @@ class SupervisedPrototypeModel(PrototypeModel):
 
     def log_acc(self, distances, targets, tag):
         preds = self.predict_from_distances(distances)
-        accuracy = torchmetrics.functional.accuracy(preds.int(), targets.int())
+        accuracy = torchmetrics.functional.accuracy(
-        # `.int()` because FloatTensors are assumed to be class probabilities
+            preds.int(),
+            targets.int(),
+            "multiclass",
+            num_classes=self.num_classes,
+        )
 
-        self.log(tag,
+        self.log(
-                 accuracy,
+            tag,
-                 on_step=False,
+            accuracy,
-                 on_epoch=True,
+            on_step=False,
-                 prog_bar=True,
+            on_epoch=True,
-                 logger=True)
+            prog_bar=True,
+            logger=True,
+        )
 
     def test_step(self, batch, batch_idx):
         x, targets = batch
 
         preds = self.predict(x)
-        accuracy = torchmetrics.functional.accuracy(preds.int(), targets.int())
+        accuracy = torchmetrics.functional.accuracy(
+            preds.int(),
+            targets.int(),
+            "multiclass",
+            num_classes=self.num_classes,
+        )
 
         self.log("test_acc", accuracy)
 
 
-class ProtoTorchMixin(object):
+class ProtoTorchMixin:
     """All mixins are ProtoTorchMixins."""
 
 
@@ -201,14 +227,16 @@ class NonGradientMixin(ProtoTorchMixin):
         super().__init__(*args, **kwargs)
         self.automatic_optimization = False
 
-    def training_step(self, train_batch, batch_idx, optimizer_idx=None):
+    def training_step(self, train_batch, batch_idx):
         raise NotImplementedError
 
 
 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)
 

src/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,
+    def __init__(
-                 threshold=0.01,
+        self,
-                 idle_epochs=10,
+        threshold=0.01,
-                 prune_quota_per_epoch=-1,
+        idle_epochs=10,
-                 frequency=1,
+        prune_quota_per_epoch=-1,
-                 replace=False,
+        frequency=1,
-                 prototypes_initializer=None,
+        replace=False,
-                 verbose=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:
+            logging.debug(f"\nPrototype win ratios: {ratios}")
-                print(f"\nPrototype win ratios: {ratios}")
+            logging.debug(f"Pruning prototypes at: {to_prune}")
-                print(f"Pruning prototypes at: {to_prune}")
+            logging.debug(f"Corresponding labels are: {prune_labels.tolist()}")
-                print(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...")
+                logging.info(f"Re-adding pruned prototypes...")
-                    print(f"distribution={distribution}")
+                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} "
+            logging.info(f"`num_prototypes` changed from {cur_num_protos} "
-                      f"to {new_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,
+                1,
-                initializer=LiteralCompInitializer(new_component.unsqueeze(0)))
+                initializer=LiteralCompInitializer(new_component.unsqueeze(0)),
+            )
 
             # Adjust Topology
             topology.add_prototype()

src/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
 
@@ -44,7 +44,7 @@ class CBC(SiameseGLVQ):
         probs = self.competition_layer(detections, reasonings)
         return probs
 
-    def shared_step(self, batch, batch_idx, optimizer_idx=None):
+    def shared_step(self, batch, batch_idx):
         x, y = batch
         y_pred = self(x)
         num_classes = self.num_classes
@@ -52,17 +52,23 @@ class CBC(SiameseGLVQ):
         loss = self.loss(y_pred, y_true).mean()
         return y_pred, loss
 
-    def training_step(self, batch, batch_idx, optimizer_idx=None):
+    def training_step(self, batch, batch_idx):
-        y_pred, train_loss = self.shared_step(batch, batch_idx, optimizer_idx)
+        y_pred, train_loss = self.shared_step(batch, batch_idx)
         preds = torch.argmax(y_pred, dim=1)
-        accuracy = torchmetrics.functional.accuracy(preds.int(),
+        accuracy = torchmetrics.functional.accuracy(
-                                                    batch[1].int())
+            preds.int(),
-        self.log("train_acc",
+            batch[1].int(),
-                 accuracy,
+            "multiclass",
-                 on_step=False,
+            num_classes=self.num_classes,
-                 on_epoch=True,
+        )
-                 prog_bar=True,
+        self.log(
-                 logger=True)
+            "train_acc",
+            accuracy,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
         return train_loss
 
     def predict(self, x):

src/prototorch/models/extras.py

@@ -5,8 +5,7 @@ Modules not yet available in prototorch go here temporarily.
 """
 
 import torch
-
+from prototorch.core.similarities import gaussian
-from ..core.similarities import gaussian
 
 
 def rank_scaled_gaussian(distances, lambd):
@@ -40,7 +39,7 @@ def ltangent_distance(x, y, omegas):
     :param `torch.tensor` omegas: Three dimensional matrix
     :rtype: `torch.tensor`
     """
-    x, y = [arr.view(arr.size(0), -1) for arr in (x, y)]
+    x, y = (arr.view(arr.size(0), -1) for arr in (x, y))
     p = torch.eye(omegas.shape[-2], device=omegas.device) - torch.bmm(
         omegas, omegas.permute([0, 2, 1]))
     projected_x = x @ p

src/prototorch/models/glvq.py

@@ -1,22 +1,22 @@
 """Models based on the GLVQ framework."""
 
 import torch
-from torch.nn.parameter import Parameter
+from prototorch.core.competitions import wtac
-
+from prototorch.core.distances import (
-from ..core.competitions import wtac
-from ..core.distances import (
     lomega_distance,
     omega_distance,
     squared_euclidean_distance,
 )
-from ..core.initializers import EyeLinearTransformInitializer
+from prototorch.core.initializers import EyeLinearTransformInitializer
-from ..core.losses import (
+from prototorch.core.losses import (
     GLVQLoss,
     lvq1_loss,
     lvq21_loss,
 )
-from ..core.transforms import LinearTransform
+from prototorch.core.transforms import LinearTransform
-from ..nn.wrappers import LambdaLayer, LossLayer
+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,
+            margin=self.hparams["margin"],
-            transfer_fn=self.hparams.transfer_fn,
+            transfer_fn=self.hparams["transfer_fn"],
-            beta=self.hparams.transfer_beta,
+            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):
@@ -66,15 +66,15 @@ class GLVQ(SupervisedPrototypeModel):
             prototype_wr,
         ])
 
-    def shared_step(self, batch, batch_idx, optimizer_idx=None):
+    def shared_step(self, batch, batch_idx):
         x, y = batch
         out = self.compute_distances(x)
         _, plabels = self.proto_layer()
         loss = self.loss(out, y, plabels)
         return out, loss
 
-    def training_step(self, batch, batch_idx, optimizer_idx=None):
+    def training_step(self, batch, batch_idx):
-        out, train_loss = self.shared_step(batch, batch_idx, optimizer_idx)
+        out, train_loss = self.shared_step(batch, batch_idx)
         self.log_prototype_win_ratios(out)
         self.log("train_loss", train_loss)
         self.log_acc(out, batch[-1], tag="train_acc")
@@ -99,10 +99,6 @@ class GLVQ(SupervisedPrototypeModel):
             test_loss += batch_loss.item()
         self.log("test_loss", test_loss)
 
-    # TODO
-    # def predict_step(self, batch, batch_idx, dataloader_idx=None):
-    #     pass
-
 
 class SiameseGLVQ(GLVQ):
     """GLVQ in a Siamese setting.
@@ -123,29 +119,9 @@ class SiameseGLVQ(GLVQ):
         self.backbone = backbone
         self.both_path_gradients = both_path_gradients
 
-    def configure_optimizers(self):
-        proto_opt = self.optimizer(self.proto_layer.parameters(),
-                                   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)
-            optimizers = [proto_opt, bb_opt]
-        else:
-            optimizers = [proto_opt]
-        if self.lr_scheduler is not None:
-            schedulers = []
-            for optimizer in optimizers:
-                scheduler = self.lr_scheduler(optimizer,
-                                              **self.lr_scheduler_kwargs)
-                schedulers.append(scheduler)
-            return optimizers, schedulers
-        else:
-            return optimizers
-
     def compute_distances(self, x):
         protos, _ = self.proto_layer()
-        x, protos = [arr.view(arr.size(0), -1) for arr in (x, protos)]
+        x, protos = (arr.view(arr.size(0), -1) for arr in (x, protos))
         latent_x = self.backbone(x)
 
         bb_grad = any([el.requires_grad for el in self.backbone.parameters()])
@@ -199,18 +175,22 @@ 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
-        self.backbone = LambdaLayer(lambda x: x @ torch.diag(self._relevances),
+        self.backbone = LambdaLayer(self._apply_relevances,
                                     name="relevance scaling")
 
+    def _apply_relevances(self, x):
+        return x @ torch.diag(self._relevances)
+
     @property
     def relevance_profile(self):
         return self._relevances.detach().cpu()
@@ -233,8 +213,8 @@ class SiameseGMLVQ(SiameseGLVQ):
         omega_initializer = kwargs.get("omega_initializer",
                                        EyeLinearTransformInitializer())
         self.backbone = LinearTransform(
-            self.hparams.input_dim,
+            self.hparams["input_dim"],
-            self.hparams.latent_dim,
+            self.hparams["latent_dim"],
             initializer=omega_initializer,
         )
 
@@ -244,7 +224,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 +237,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,11 +247,9 @@ class GMLVQ(GLVQ):
         # Additional parameters
         omega_initializer = kwargs.get("omega_initializer",
                                        EyeLinearTransformInitializer())
-        omega = omega_initializer.generate(self.hparams.input_dim,
+        omega = omega_initializer.generate(self.hparams["input_dim"],
-                                           self.hparams.latent_dim)
+                                           self.hparams["latent_dim"])
         self.register_parameter("_omega", Parameter(omega))
-        self.backbone = LambdaLayer(lambda x: x @ self._omega,
-                                    name="omega matrix")
 
     @property
     def omega_matrix(self):
@@ -299,8 +280,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["input_dim"],
-            self.hparams.latent_dim,
+            self.hparams["latent_dim"],
             device=self.device,
         )
         self.register_parameter("_omega", Parameter(omega))
@@ -316,23 +297,27 @@ class GTLVQ(LGMLVQ):
         omega_initializer = kwargs.get("omega_initializer")
 
         if omega_initializer is not None:
-            subspace = omega_initializer.generate(self.hparams.input_dim,
+            subspace = omega_initializer.generate(
-                                                  self.hparams.latent_dim)
+                self.hparams["input_dim"],
-            omega = torch.repeat_interleave(subspace.unsqueeze(0),
+                self.hparams["latent_dim"],
-                                            self.num_prototypes,
+            )
-                                            dim=0)
+            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["input_dim"],
-                self.hparams.latent_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 +374,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():

src/prototorch/models/knn.py

@@ -2,13 +2,14 @@
 
 import warnings
 
-from ..core.competitions import KNNC
+from prototorch.core.competitions import KNNC
-from ..core.components import LabeledComponents
+from prototorch.core.components import LabeledComponents
-from ..core.initializers import (
+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
 
 
@@ -33,13 +34,10 @@ class KNN(SupervisedPrototypeModel):
             labels_initializer=LiteralLabelsInitializer(targets))
         self.competition_layer = KNNC(k=self.hparams.k)
 
-    def training_step(self, train_batch, batch_idx, optimizer_idx=None):
+    def training_step(self, train_batch, batch_idx):
         return 1  # skip training step
 
-    def on_train_batch_start(self,
+    def on_train_batch_start(self, train_batch, batch_idx):
-                             train_batch,
-                             batch_idx,
-                             dataloader_idx=None):
         warnings.warn("k-NN has no training, skipping!")
         return -1
 

src/prototorch/models/lvq.py

@@ -1,8 +1,11 @@
 """LVQ models that are optimized using non-gradient methods."""
 
-from ..core.losses import _get_dp_dm
+import logging
-from ..nn.activations import get_activation
+
-from ..nn.wrappers import LambdaLayer
+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
 
@@ -10,7 +13,7 @@ from .glvq import GLVQ
 class LVQ1(NonGradientMixin, GLVQ):
     """Learning Vector Quantization 1."""
 
-    def training_step(self, train_batch, batch_idx, optimizer_idx=None):
+    def training_step(self, train_batch, batch_idx):
         protos, plables = self.proto_layer()
         x, y = train_batch
         dis = self.compute_distances(x)
@@ -29,8 +32,8 @@ class LVQ1(NonGradientMixin, GLVQ):
             self.proto_layer.load_state_dict({"_components": updated_protos},
                                              strict=False)
 
-        print(f"dis={dis}")
+        logging.debug(f"dis={dis}")
-        print(f"y={y}")
+        logging.debug(f"y={y}")
         # Logging
         self.log_acc(dis, y, tag="train_acc")
 
@@ -40,7 +43,7 @@ class LVQ1(NonGradientMixin, GLVQ):
 class LVQ21(NonGradientMixin, GLVQ):
     """Learning Vector Quantization 2.1."""
 
-    def training_step(self, train_batch, batch_idx, optimizer_idx=None):
+    def training_step(self, train_batch, batch_idx):
         protos, plabels = self.proto_layer()
 
         x, y = train_batch
@@ -73,8 +76,7 @@ class MedianLVQ(NonGradientMixin, GLVQ):
 
     """
 
-    def __init__(self, hparams, verbose=True, **kwargs):
+    def __init__(self, hparams, **kwargs):
-        self.verbose = verbose
         super().__init__(hparams, **kwargs)
 
         self.transfer_layer = LambdaLayer(
@@ -98,7 +100,7 @@ class MedianLVQ(NonGradientMixin, GLVQ):
         lower_bound = (gamma * f.log()).sum()
         return lower_bound
 
-    def training_step(self, train_batch, batch_idx, optimizer_idx=None):
+    def training_step(self, train_batch, batch_idx):
         protos, plabels = self.proto_layer()
 
         x, y = train_batch
@@ -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:
+                    logging.debug(f"Updating prototype {i} to data {k}...")
-                        print(f"Updating prototype {i} to data {k}...")
                     self.proto_layer.load_state_dict({"_components": _protos},
                                                      strict=False)
                     break

src/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
 
@@ -18,7 +21,7 @@ class CELVQ(GLVQ):
         # Loss
         self.loss = torch.nn.CrossEntropyLoss()
 
-    def shared_step(self, batch, batch_idx, optimizer_idx=None):
+    def shared_step(self, batch, batch_idx):
         x, y = batch
         out = self.compute_distances(x)  # [None, num_protos]
         _, plabels = self.proto_layer()
@@ -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):
@@ -53,7 +63,7 @@ class ProbabilisticLVQ(GLVQ):
         prediction[confidence < self.rejection_confidence] = -1
         return prediction
 
-    def training_step(self, batch, batch_idx, optimizer_idx=None):
+    def training_step(self, batch, batch_idx):
         x, y = batch
         out = self.forward(x)
         _, plabels = self.proto_layer()
@@ -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)
 
 
@@ -107,7 +123,7 @@ class PLVQ(ProbabilisticLVQ, SiameseGMLVQ):
         self.loss = torch.nn.KLDivLoss()
 
     # FIXME
-    # def training_step(self, batch, batch_idx, optimizer_idx=None):
+    # def training_step(self, batch, batch_idx):
     #     x, y = batch
     #     y_pred = self(x)
     #     batch_loss = self.loss(y_pred, y)

src/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")
@@ -63,7 +63,7 @@ class KohonenSOM(NonGradientMixin, UnsupervisedPrototypeModel):
             strict=False,
         )
 
-    def training_epoch_end(self, training_step_outputs):
+    def on_training_epoch_end(self, training_step_outputs):
         self._sigma = self.hparams.sigma * np.exp(
             -self.current_epoch / self.trainer.max_epochs)
 
@@ -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,
+            agelimit=self.hparams["age_limit"],
-            num_prototypes=self.hparams.num_prototypes,
+            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,
+                reduction=self.hparams["insert_reduction"],
-                freq=self.hparams.insert_freq,
+                freq=self.hparams["insert_freq"],
             )
         ]

src/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))))
+                if isinstance(data, Sized):
-            elif isinstance(data, torch.utils.data.DataLoader):
+                    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):
+    def on_train_start(self, _, pl_module):
-        tb = pl_module.logger.experiment
+        if isinstance(pl_module.logger, TensorBoardLogger):
-        if self.add_embedding:
+            tb = pl_module.logger.experiment
-            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)
 
-        if self.random_data:
+            # Add embedding
-            ind = np.random.choice(len(self.x_train),
+            if self.add_embedding:
-                                   size=self.random_data,
+                if self.x_train is not None and self.y_train is not None:
-                                   replace=False)
+                    ind = np.random.choice(len(self.x_train),
-            data = self.x_train[ind]
+                                           size=self.embedding_data,
-            grid = torchvision.utils.make_grid(data, nrow=self.num_columns)
+                                           replace=False)
-            tb.add_image(tag="Data",
+                    data = self.x_train[ind]
-                         img_tensor=grid,
+                    tb.add_embedding(data.view(len(ind), -1),
-                         global_step=None,
+                                     label_img=data,
-                         dataformats=self.dataformats)
+                                     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

tests/test_models.py

@@ -1,195 +1,193 @@
 """prototorch.models test suite."""
 
-import prototorch as pt
+import prototorch.models
-import pytest
-import torch
 
 
 def test_glvq_model_build():
-    model = pt.models.GLVQ(
+    model = prototorch.models.GLVQ(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_glvq1_model_build():
-    model = pt.models.GLVQ1(
+    model = prototorch.models.GLVQ1(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_glvq21_model_build():
-    model = pt.models.GLVQ1(
+    model = prototorch.models.GLVQ1(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_gmlvq_model_build():
-    model = pt.models.GMLVQ(
+    model = prototorch.models.GMLVQ(
         {
             "distribution": (3, 2),
             "input_dim": 2,
             "latent_dim": 2,
         },
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_grlvq_model_build():
-    model = pt.models.GRLVQ(
+    model = prototorch.models.GRLVQ(
         {
             "distribution": (3, 2),
             "input_dim": 2,
         },
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_gtlvq_model_build():
-    model = pt.models.GTLVQ(
+    model = prototorch.models.GTLVQ(
         {
             "distribution": (3, 2),
             "input_dim": 4,
             "latent_dim": 2,
         },
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_lgmlvq_model_build():
-    model = pt.models.LGMLVQ(
+    model = prototorch.models.LGMLVQ(
         {
             "distribution": (3, 2),
             "input_dim": 4,
             "latent_dim": 2,
         },
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_image_glvq_model_build():
-    model = pt.models.ImageGLVQ(
+    model = prototorch.models.ImageGLVQ(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(16),
+        prototypes_initializer=prototorch.initializers.RNCI(16),
     )
 
 
 def test_image_gmlvq_model_build():
-    model = pt.models.ImageGMLVQ(
+    model = prototorch.models.ImageGMLVQ(
         {
             "distribution": (3, 2),
             "input_dim": 16,
             "latent_dim": 2,
         },
-        prototypes_initializer=pt.initializers.RNCI(16),
+        prototypes_initializer=prototorch.initializers.RNCI(16),
     )
 
 
 def test_image_gtlvq_model_build():
-    model = pt.models.ImageGMLVQ(
+    model = prototorch.models.ImageGMLVQ(
         {
             "distribution": (3, 2),
             "input_dim": 16,
             "latent_dim": 2,
         },
-        prototypes_initializer=pt.initializers.RNCI(16),
+        prototypes_initializer=prototorch.initializers.RNCI(16),
     )
 
 
 def test_siamese_glvq_model_build():
-    model = pt.models.SiameseGLVQ(
+    model = prototorch.models.SiameseGLVQ(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(4),
+        prototypes_initializer=prototorch.initializers.RNCI(4),
     )
 
 
 def test_siamese_gmlvq_model_build():
-    model = pt.models.SiameseGMLVQ(
+    model = prototorch.models.SiameseGMLVQ(
         {
             "distribution": (3, 2),
             "input_dim": 4,
             "latent_dim": 2,
         },
-        prototypes_initializer=pt.initializers.RNCI(4),
+        prototypes_initializer=prototorch.initializers.RNCI(4),
     )
 
 
 def test_siamese_gtlvq_model_build():
-    model = pt.models.SiameseGTLVQ(
+    model = prototorch.models.SiameseGTLVQ(
         {
             "distribution": (3, 2),
             "input_dim": 4,
             "latent_dim": 2,
         },
-        prototypes_initializer=pt.initializers.RNCI(4),
+        prototypes_initializer=prototorch.initializers.RNCI(4),
     )
 
 
 def test_knn_model_build():
-    train_ds = pt.datasets.Iris(dims=[0, 2])
+    train_ds = prototorch.datasets.Iris(dims=[0, 2])
-    model = pt.models.KNN(dict(k=3), data=train_ds)
+    model = prototorch.models.KNN(dict(k=3), data=train_ds)
 
 
 def test_lvq1_model_build():
-    model = pt.models.LVQ1(
+    model = prototorch.models.LVQ1(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_lvq21_model_build():
-    model = pt.models.LVQ21(
+    model = prototorch.models.LVQ21(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_median_lvq_model_build():
-    model = pt.models.MedianLVQ(
+    model = prototorch.models.MedianLVQ(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_celvq_model_build():
-    model = pt.models.CELVQ(
+    model = prototorch.models.CELVQ(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_rslvq_model_build():
-    model = pt.models.RSLVQ(
+    model = prototorch.models.RSLVQ(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_slvq_model_build():
-    model = pt.models.SLVQ(
+    model = prototorch.models.SLVQ(
         {"distribution": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_growing_neural_gas_model_build():
-    model = pt.models.GrowingNeuralGas(
+    model = prototorch.models.GrowingNeuralGas(
         {"num_prototypes": 5},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_kohonen_som_model_build():
-    model = pt.models.KohonenSOM(
+    model = prototorch.models.KohonenSOM(
         {"shape": (3, 2)},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )
 
 
 def test_neural_gas_model_build():
-    model = pt.models.NeuralGas(
+    model = prototorch.models.NeuralGas(
         {"num_prototypes": 5},
-        prototypes_initializer=pt.initializers.RNCI(2),
+        prototypes_initializer=prototorch.initializers.RNCI(2),
     )