build: bump version 0.7.0 → 0.7.1

Necessary for different optimizer runs

.bumpversion.cfg

@@ -1,13 +1,13 @@
 [bumpversion]
-current_version = 0.4.1
+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,16 +6,16 @@ name: examples
 on:
   push:
     paths:
-      - 'examples/**.py'
+      - "examples/**.py"
 jobs:
   cpu:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.9
-      uses: actions/setup-python@v2
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v4
         with:
-        python-version: 3.9
+          python-version: "3.11"
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip

.github/workflows/pythonapp.yml

@@ -6,40 +6,42 @@ name: tests
 on:
   push:
   pull_request:
-    branches: [ master ]
+    branches: [master]
 
 jobs:
   style:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.9
-      uses: actions/setup-python@v2
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v4
         with:
-        python-version: 3.9
+          python-version: "3.11"
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           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"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         os: [ubuntu-latest, windows-latest]
         exclude:
-        - os: windows-latest
-          python-version: "3.7"
           - os: windows-latest
             python-version: "3.8"
+          - os: windows-latest
+            python-version: "3.9"
+          - os: windows-latest
+            python-version: "3.10"
 
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
@@ -54,18 +56,18 @@ jobs:
     needs: compatibility
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.9
-      uses: actions/setup-python@v2
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v4
         with:
-        python-version: "3.9"
+          python-version: "3.11"
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           pip install .[all]
-        pip install wheel
+          pip install build
       - name: Build package
-      run: python setup.py sdist bdist_wheel
+        run: python -m build . -C verbose
       - name: Publish a Python distribution to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:

.pre-commit-config.yaml

@@ -2,8 +2,8 @@
 # See https://pre-commit.com/hooks.html for more hooks
 
 repos:
-- repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.1.0
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
@@ -12,41 +12,42 @@ repos:
       - id: check-ast
       - id: check-case-conflict
 
-- repo: https://github.com/myint/autoflake
-  rev: v1.4
+  - repo: https://github.com/myint/autoflake
+    rev: v2.1.1
     hooks:
       - id: autoflake
 
-- repo: http://github.com/PyCQA/isort
-  rev: 5.10.1
+  - repo: http://github.com/PyCQA/isort
+    rev: 5.12.0
     hooks:
       - id: isort
 
-- repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.931
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.3.0
     hooks:
       - id: mypy
         files: prototorch
         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
     hooks:
       - id: yapf
+        additional_dependencies: ["toml"]
 
-- repo: https://github.com/pre-commit/pygrep-hooks
-  rev: v1.9.0
+  - repo: https://github.com/pre-commit/pygrep-hooks
+    rev: v1.10.0
     hooks:
       - id: python-use-type-annotations
       - id: python-no-log-warn
       - id: python-check-blanket-noqa
 
-- repo: https://github.com/asottile/pyupgrade
-  rev: v2.31.0
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v3.7.0
     hooks:
       - id: pyupgrade
 
-- repo: https://github.com/si-cim/gitlint
+  - repo: https://github.com/si-cim/gitlint
     rev: v0.15.2-unofficial
     hooks:
       - id: gitlint

docs/source/conf.py

@@ -23,7 +23,7 @@ author = "Jensun Ravichandran"
 
 # The full version, including alpha/beta/rc tags
 #
-release = "0.4.1"
+release = "0.7.1"
 
 # -- General configuration ---------------------------------------------------
 

docs/source/tutorial.ipynb

@@ -2,223 +2,252 @@
  "cells": [
   {
    "cell_type": "markdown",
+   "id": "7ac5eff0",
+   "metadata": {},
    "source": [
     "# A short tutorial for the `prototorch.models` plugin"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "beb83780",
+   "metadata": {},
    "source": [
     "## Introduction"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "43b74278",
+   "metadata": {},
    "source": [
-    "This is a short tutorial for the [models](https://github.com/si-cim/prototorch_models) plugin of the [ProtoTorch](https://github.com/si-cim/prototorch) framework.\n",
+    "This is a short tutorial for the [models](https://github.com/si-cim/prototorch_models) plugin of the [ProtoTorch](https://github.com/si-cim/prototorch) framework. This is by no means a comprehensive look at all the features that the framework has to offer, but it should help you get up and running.\n",
     "\n",
     "[ProtoTorch](https://github.com/si-cim/prototorch) provides [torch.nn](https://pytorch.org/docs/stable/nn.html) modules and utilities to implement prototype-based models. However, it is up to the user to put these modules together into models and handle the training of these models. Expert machine-learning practioners and researchers sometimes prefer this level of control. However, this leads to a lot of boilerplate code that is essentially same across many projects. Needless to say, this is a source of a lot of frustration. [PyTorch-Lightning](https://pytorch-lightning.readthedocs.io/en/latest/) is a framework that helps avoid a lot of this frustration by handling the boilerplate code for you so you don't have to reinvent the wheel every time you need to implement a new model.\n",
     "\n",
     "With the [prototorch.models](https://github.com/si-cim/prototorch_models) plugin, we've gone one step further and pre-packaged commonly used prototype-models like GMLVQ as [Lightning-Modules](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.core.lightning.html?highlight=lightning%20module#pytorch_lightning.core.lightning.LightningModule). With only a few lines to code, it is now possible to build and train prototype-models. It quite simply cannot get any simpler than this."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "4e5d1fad",
+   "metadata": {},
    "source": [
     "## Basics"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "1244b66b",
+   "metadata": {},
    "source": [
     "First things first. When working with the models plugin, you'll probably need `torch`, `prototorch` and `pytorch_lightning`. So, we recommend that you import all three like so:"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "dcb88e8a",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "import prototorch as pt\n",
     "import pytorch_lightning as pl\n",
     "import torch"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "1adbe2f8",
+   "metadata": {},
    "source": [
     "### Building Models"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "96663ab1",
+   "metadata": {},
    "source": [
     "Let's start by building a `GLVQ` model. It is one of the simplest models to build. The only requirements are a prototype distribution and an initializer."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "819ba756",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "model = pt.models.GLVQ(\n",
     "    hparams=dict(distribution=[1, 1, 1]),\n",
     "    prototypes_initializer=pt.initializers.ZerosCompInitializer(2),\n",
     ")"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "1b37e97c",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "print(model)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "d2c86903",
+   "metadata": {},
    "source": [
     "The key `distribution` in the `hparams` argument describes the prototype distribution. If it is a Python [list](https://docs.python.org/3/tutorial/datastructures.html), it is assumed that there are as many entries in this list as there are classes, and the number at each location of this list describes the number of prototypes to be used for that particular class. So, `[1, 1, 1]` implies that we have three classes with one prototype per class.   If it is a Python [tuple](https://docs.python.org/3/tutorial/datastructures.html), a shorthand of `(num_classes, prototypes_per_class)` is assumed. If it is a Python [dictionary](https://docs.python.org/3/tutorial/datastructures.html), the key-value pairs describe the class label and the number of prototypes for that class respectively. So, `{0: 2, 1: 2, 2: 2}` implies that we have three classes with labels `{1, 2, 3}`, each equipped with two prototypes. If however, the dictionary contains the keys `\"num_classes\"` and `\"per_class\"`, they are parsed to use their values as one might expect.\n",
     "\n",
     "The `prototypes_initializer` argument describes how the prototypes are meant to be initialized. This argument has to be an instantiated object of some kind of [AbstractComponentsInitializer](https://github.com/si-cim/prototorch/blob/dev/prototorch/components/initializers.py#L18). If this is a [ShapeAwareCompInitializer](https://github.com/si-cim/prototorch/blob/dev/prototorch/components/initializers.py#L41), this only requires a `shape` arugment that describes the shape of the prototypes. So, `pt.initializers.ZerosCompInitializer(3)` creates 3d-vector prototypes all initialized to zeros."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "45806052",
+   "metadata": {},
    "source": [
     "### Data"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "9d62c4c6",
+   "metadata": {},
    "source": [
     "The preferred way to working with data in `torch` is to use the [Dataset and Dataloader API](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html). There a few pre-packaged datasets available under `prototorch.datasets`. See [here](https://prototorch.readthedocs.io/en/latest/api.html#module-prototorch.datasets) for a full list of available datasets."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "504df02c",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_ds = pt.datasets.Iris(dims=[0, 2])"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "3b8e7756",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "type(train_ds)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "bce43afa",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_ds.data.shape, train_ds.targets.shape"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "26a83328",
+   "metadata": {},
    "source": [
     "Once we have such a dataset, we could wrap it in a `Dataloader` to load the data in batches, and possibly apply some transformations on the fly."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "67b80fbe",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_loader = torch.utils.data.DataLoader(train_ds, batch_size=2)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "c1185f31",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "type(train_loader)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "9b5a8963",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "x_batch, y_batch = next(iter(train_loader))\n",
     "print(f\"{x_batch=}, {y_batch=}\")"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "dd492ee2",
+   "metadata": {},
    "source": [
     "This perhaps seems like a lot of work for a small dataset that fits completely in memory. However, this comes in very handy when dealing with huge datasets that can only be processed in batches."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "5176b055",
+   "metadata": {},
    "source": [
     "### Training"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "46a7a506",
+   "metadata": {},
    "source": [
     "If you're familiar with other deep learning frameworks, you might perhaps expect a `.fit(...)` or `.train(...)` method. However, in PyTorch-Lightning, this is done slightly differently. We first create a trainer and then pass the model and the Dataloader to `trainer.fit(...)` instead. So, it is more functional in style than object-oriented."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "279e75b7",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "trainer = pl.Trainer(max_epochs=2, weights_summary=None)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "e496b492",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "trainer.fit(model, train_loader)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "497fbff6",
+   "metadata": {},
    "source": [
     "### From data to a trained model - a very minimal example"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "ab069c5d",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_ds = pt.datasets.Iris(dims=[0, 2])\n",
     "train_loader = torch.utils.data.DataLoader(train_ds, batch_size=32)\n",
@@ -230,49 +259,239 @@
     "\n",
     "trainer = pl.Trainer(max_epochs=50, weights_summary=None)\n",
     "trainer.fit(model, train_loader)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "30c71a93",
+   "metadata": {},
    "source": [
-    "## Advanced"
-   ],
-   "metadata": {}
+    "### Saving/Loading trained models"
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "f74ed2c1",
+   "metadata": {},
    "source": [
-    "### Initializing prototypes with a subset of a dataset (along with transformations)"
-   ],
-   "metadata": {}
+    "Pytorch Lightning can automatically checkpoint the model during various stages of training, but it also possible to manually save a checkpoint after training."
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "3156658d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ckpt_path = \"./checkpoints/glvq_iris.ckpt\"\n",
+    "trainer.save_checkpoint(ckpt_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c1c34055",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loaded_model = pt.models.GLVQ.load_from_checkpoint(ckpt_path, strict=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bbbb08e9",
+   "metadata": {},
+   "source": [
+    "### Visualizing decision boundaries in 2D"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53ca52dc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pt.models.VisGLVQ2D(data=train_ds).visualize(loaded_model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8373531f",
+   "metadata": {},
+   "source": [
+    "### Saving/Loading trained weights"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "937bc458",
+   "metadata": {},
+   "source": [
+    "In most cases, the checkpointing workflow is sufficient. In some cases however, one might want to only save the trained weights from the model. The disadvantage of this method is that the model has be re-created using compatible initialization parameters before the weights could be loaded."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1f2035af",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ckpt_path = \"./checkpoints/glvq_iris_weights.pth\"\n",
+    "torch.save(model.state_dict(), ckpt_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1206021a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = pt.models.GLVQ(\n",
+    "    dict(distribution=(3, 2)),\n",
+    "    prototypes_initializer=pt.initializers.ZerosCompInitializer(2),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9f2a4beb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pt.models.VisGLVQ2D(data=train_ds, title=\"Before loading the weights\").visualize(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "528d2fc2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch.load(ckpt_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ec817e6b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.load_state_dict(torch.load(ckpt_path), strict=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a208eab7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pt.models.VisGLVQ2D(data=train_ds, title=\"After loading the weights\").visualize(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f8de748f",
+   "metadata": {},
+   "source": [
+    "## Advanced"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "53a64063",
+   "metadata": {},
+   "source": [
+    "### Warm-start a model with prototypes learned from another model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3177c277",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trained_model = pt.models.GLVQ.load_from_checkpoint(\"./checkpoints/glvq_iris.ckpt\", strict=False)\n",
+    "model = pt.models.SiameseGMLVQ(\n",
+    "    dict(input_dim=2,\n",
+    "         latent_dim=2,\n",
+    "         distribution=(3, 2),\n",
+    "         proto_lr=0.0001,\n",
+    "         bb_lr=0.0001),\n",
+    "    optimizer=torch.optim.Adam,\n",
+    "    prototypes_initializer=pt.initializers.LCI(trained_model.prototypes),\n",
+    "    labels_initializer=pt.initializers.LLI(trained_model.prototype_labels),\n",
+    "    omega_initializer=pt.initializers.LLTI(torch.tensor([[0., 1.], [1., 0.]])),  # permute axes\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8baee9a2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cc203088",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pt.models.VisSiameseGLVQ2D(data=train_ds, title=\"GMLVQ - Warm-start state\").visualize(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1f6a33a5",
+   "metadata": {},
+   "source": [
+    "### Initializing prototypes with a subset of a dataset (along with transformations)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "946ce341",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "import prototorch as pt\n",
     "import pytorch_lightning as pl\n",
     "import torch\n",
     "from torchvision import transforms\n",
-    "from torchvision.datasets import MNIST"
-   ],
-   "outputs": [],
-   "metadata": {}
+    "from torchvision.datasets import MNIST\n",
+    "from torchvision.utils import make_grid"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "510d9bd4",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "from matplotlib import pyplot as plt"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "ea7c1228",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_ds = MNIST(\n",
     "    \"~/datasets\",\n",
@@ -284,59 +503,87 @@
     "        transforms.ToTensor(),\n",
     "    ]),\n",
     ")"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "1b9eaf5c",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "s = int(0.05 * len(train_ds))\n",
     "init_ds, rest_ds = torch.utils.data.random_split(train_ds, [s, len(train_ds) - s])"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "8c32c9f2",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "init_ds"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "68a9a8b9",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "model = pt.models.ImageGLVQ(\n",
-    "    dict(distribution=(10, 5)),\n",
+    "    dict(distribution=(10, 1)),\n",
     "    prototypes_initializer=pt.initializers.SMCI(init_ds),\n",
     ")"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "source": [
-    "plt.imshow(model.get_prototype_grid(num_columns=10))"
-   ],
+   "id": "6f23df86",
+   "metadata": {},
    "outputs": [],
-   "metadata": {}
+   "source": [
+    "plt.imshow(model.get_prototype_grid(num_columns=5))"
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "1c23c7b2",
+   "metadata": {},
+   "source": [
+    "We could, of course, just use the initializers in isolation. For example, we could quickly obtain a stratified selection from the data like so:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "30780927",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "protos, plabels = pt.components.LabeledComponents(\n",
+    "    distribution=(10, 5),\n",
+    "    components_initializer=pt.initializers.SMCI(init_ds),\n",
+    "    labels_initializer=pt.initializers.LabelsInitializer(),\n",
+    ")()\n",
+    "plt.imshow(make_grid(protos, 10).permute(1, 2, 0)[:, :, 0], cmap=\"jet\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4fa69f92",
+   "metadata": {},
    "source": [
     "## FAQs"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "fa20f9ac",
+   "metadata": {},
    "source": [
     "### How do I Retrieve the prototypes and their respective labels from the model?\n",
     "\n",
@@ -351,11 +598,12 @@
     "```python\n",
     ">>> model.prototype_labels\n",
     "```"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "ba8215bf",
+   "metadata": {},
    "source": [
     "### How do I make inferences/predictions/recall with my trained model?\n",
     "\n",
@@ -370,13 +618,12 @@
     "```python\n",
     ">>> y_pred = model(torch.Tensor(x_train))  # returns probabilities\n",
     "```"
-   ],
-   "metadata": {}
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -390,7 +637,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.4"
+   "version": "3.9.12"
   }
  },
  "nbformat": 4,

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),
-        reasonings_iniitializer=pt.initializers.
+        components_initializer=pt.initializers.SSCI(train_ds, noise=0.1),
+        reasonings_initializer=pt.initializers.
         PurePositiveReasoningsInitializer(),
     )
 
     # Callbacks
-    vis = pt.models.VisCBC2D(data=train_ds,
+    vis = VisCBC2D(
+        data=train_ds,
         title="CBC Iris Example",
         resolution=100,
-                             axis_off=True)
+        axis_off=True,
+    )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
-        args,
-        callbacks=[vis],
+    trainer = pl.Trainer(
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
+        callbacks=[
+            vis,
+        ],
+        detect_anomaly=True,
+        log_every_n_steps=1,
+        max_epochs=1000,
     )
 
     # 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_samples=500,
         num_classes=num_classes,
         num_features=num_features,
         num_clusters=num_clusters,
         separation=3.0,
-                                  seed=42)
+        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)
-    pruning = pt.models.PruneLoserPrototypes(
+    vis = VisGLVQ2D(train_ds)
+    pruning = PruneLoserPrototypes(
         threshold=0.01,  # prune prototype if it wins less than 1%
         idle_epochs=20,  # pruning too early may cause problems
         prune_quota_per_epoch=2,  # prune at most 2 prototypes per epoch
         frequency=1,  # prune every epoch
         verbose=True,
     )
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
         monitor="train_loss",
         min_delta=0.001,
         patience=20,
@@ -64,17 +84,18 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
-        args,
+    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,
             pruning,
             es,
         ],
-        progress_bar_refresh_rate=0,
-        terminate_on_nan=True,
-        weights_summary="full",
-        accelerator="ddp",
+        detect_anomaly=True,
+        log_every_n_steps=1,
+        max_epochs=1000,
     )
 
     # Training loop

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,15 +53,30 @@ if __name__ == "__main__":
     model.example_input_array = torch.zeros(4, 2)
 
     # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
+    vis = VisGLVQ2D(data=train_ds)
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
-        args,
-        callbacks=[vis],
-        weights_summary="full",
-        accelerator="ddp",
+    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=100,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop
     trainer.fit(model, train_loader)
+
+    # Manual save
+    trainer.save_checkpoint("./glvq_iris.ckpt")
+
+    # Load saved model
+    new_model = GLVQ.load_from_checkpoint(
+        checkpoint_path="./glvq_iris.ckpt",
+        strict=False,
+    )
+    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(
-        args,
-        callbacks=[vis],
-        weights_summary="full",
-        accelerator="ddp",
+    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=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,
-                                               num_workers=0,
-                                               batch_size=256)
-    test_loader = torch.utils.data.DataLoader(test_ds,
-                                              num_workers=0,
-                                              batch_size=256)
+    train_loader = DataLoader(train_ds, num_workers=4, batch_size=256)
+    test_loader = DataLoader(test_ds, num_workers=4, batch_size=256)
 
     # Hyperparameters
     num_classes = 10
@@ -52,14 +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(
-        args,
+    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,
             pruning,
-            # es,
+            es,
         ],
-        terminate_on_nan=True,
-        weights_summary=None,
-        # accelerator="ddp",
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

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(
-        args,
+    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,
             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(
-        args,
+    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=100,
-        callbacks=[vis],
-        weights_summary="full",
+        log_every_n_steps=1,
+        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,
-                                               num_workers=0,
-                                               batch_size=256)
-    test_loader = torch.utils.data.DataLoader(test_ds,
-                                              num_workers=0,
-                                              batch_size=256)
+    train_loader = DataLoader(train_ds, num_workers=0, batch_size=256)
+    test_loader = DataLoader(test_ds, num_workers=0, batch_size=256)
 
     # Hyperparameters
     num_classes = 10
@@ -52,7 +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(
-        args,
+    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,
             pruning,
-            # es,
+            es,
         ],
-        terminate_on_nan=True,
-        weights_summary=None,
-        accelerator="ddp",
+        max_epochs=1000,
+        log_every_n_steps=1,
+        detect_anomaly=True,
     )
 
     # Training loop

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(
+        train_ds,
         batch_size=256,
-                                               shuffle=True)
+        shuffle=True,
+    )
 
     # Hyperparameters
     # Latent_dim should be lower than input dim.
     hparams = dict(distribution=[1, 3], input_dim=2, latent_dim=1)
 
     # Initialize the model
-    model = pt.models.GTLVQ(
-        hparams, prototypes_initializer=pt.initializers.SMCI(train_ds))
+    model = GTLVQ(hparams,
+                  prototypes_initializer=pt.initializers.SMCI(train_ds))
 
     # Compute intermediate input and output sizes
     model.example_input_array = torch.zeros(4, 2)
 
     # Summary
-    print(model)
+    logging.info(model)
 
     # Callbacks
-    vis = pt.models.VisGLVQ2D(data=train_ds)
-    es = pl.callbacks.EarlyStopping(
+    vis = VisGLVQ2D(data=train_ds)
+    es = EarlyStopping(
         monitor="train_acc",
         min_delta=0.001,
         patience=20,
@@ -49,14 +62,17 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
-        args,
+    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,
             es,
         ],
-        weights_summary="full",
-        accelerator="ddp",
+        max_epochs=1000,
+        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(
+        X,
         y,
         test_size=0.5,
-                                                        random_state=42)
+        random_state=42,
+    )
 
     train_ds = pt.datasets.NumpyDataset(X_train, y_train)
     test_ds = pt.datasets.NumpyDataset(X_test, y_test)
 
     # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=16)
-    test_loader = torch.utils.data.DataLoader(test_ds, batch_size=16)
+    train_loader = DataLoader(train_ds, batch_size=16)
+    test_loader = DataLoader(test_ds, batch_size=16)
 
     # Hyperparameters
     hparams = dict(k=5)
 
     # Initialize the model
-    model = pt.models.KNN(hparams, data=train_ds)
+    model = KNN(hparams, data=train_ds)
 
     # Compute intermediate input and output sizes
     model.example_input_array = torch.zeros(4, 2)
 
     # Summary
-    print(model)
+    logging.info(model)
 
     # Callbacks
-    vis = pt.models.VisGLVQ2D(
+    vis = VisGLVQ2D(
         data=(X_train, y_train),
         resolution=200,
         block=True,
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
-        args,
+    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,
         max_epochs=1,
-        callbacks=[vis],
-        weights_summary="full",
+        callbacks=[
+            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[1],
                 iloc[0],
-                     cnames[i],
+                color_names[i],
                 ha="center",
                 va="center",
-                     bbox=dict(facecolor="white", alpha=0.5, lw=0))
+                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_loader = torch.utils.data.DataLoader(train_ds, batch_size=8)
+    train_ds = TensorDataset(data)
+    train_loader = DataLoader(train_ds, batch_size=8)
 
     # Hyperparameters
     hparams = dict(
@@ -78,7 +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(
-        args,
+    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,
         max_epochs=500,
-        callbacks=[vis],
-        weights_summary="full",
+        callbacks=[
+            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,
-                                               batch_size=256,
-                                               shuffle=True)
+    train_loader = DataLoader(train_ds, 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)
-    es = pl.callbacks.EarlyStopping(
+    vis = VisGLVQ2D(data=train_ds)
+    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(
-        args,
+    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,
             es,
         ],
-        weights_summary="full",
-        accelerator="ddp",
+        log_every_n_steps=1,
+        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(
-        args,
+    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,
             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)
-    es = pl.callbacks.EarlyStopping(
+    vis = VisGLVQ2D(data=train_ds)
+    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(
-        args,
-        callbacks=[vis, es],
-        weights_summary="full",
+    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,
+            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(
-        args,
-        callbacks=[vis],
-        weights_summary="full",
+    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=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(
-        args,
-        callbacks=[vis],
-        terminate_on_nan=True,
-        weights_summary="full",
-        accelerator="ddp",
+    trainer = pl.Trainer(
+        accelerator="cuda" if args.gpus else "cpu",
+        devices=args.gpus if args.gpus else "auto",
+        fast_dev_run=args.fast_dev_run,
+        callbacks=[
+            vis,
+        ],
+        detect_anomaly=True,
+        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,
-        bb_lr=0.01,
+        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(
-        args,
-        callbacks=[vis],
+    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=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],
-                   proto_lr=0.01,
-                   bb_lr=0.01,
+    hparams = dict(
+        distribution=[1, 2, 3],
+        lr=0.01,
         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(
-        args,
-        callbacks=[vis],
+    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=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,
-        callbacks=[es],
-        weights_summary=None,
+        accelerator="cpu",
+        max_epochs=50 if args.fast_dev_run else
+        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)
-    pruning = pt.models.PruneLoserPrototypes(
+    vis = VisGLVQ2D(data=train_ds)
+    pruning = PruneLoserPrototypes(
         threshold=0.02,
         idle_epochs=2,
         prune_quota_per_epoch=5,
         frequency=1,
         verbose=True,
     )
-    es = pl.callbacks.EarlyStopping(
+    es = EarlyStopping(
         monitor="train_loss",
         min_delta=0.001,
         patience=10,
@@ -88,15 +111,18 @@ if __name__ == "__main__":
     )
 
     # Setup trainer
-    trainer = pl.Trainer.from_argparse_args(
-        args,
+    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,
             pruning,
             es,
         ],
-        weights_summary="full",
-        accelerator="ddp",
+        max_epochs=1000,
+        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.0",
-    "pytorch_lightning>=1.3.5",
-    "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.4.1",
-    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.6",
-    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.9",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.7",
-        "Programming Language :: Python :: 3.6",
-        "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

@@ -26,16 +26,14 @@ from .lvq import (
 )
 from .probabilistic import (
     CELVQ,
-    PLVQ,
     RSLVQ,
     SLVQ,
 )
 from .unsupervised import (
     GrowingNeuralGas,
-    HeskesSOM,
     KohonenSOM,
     NeuralGas,
 )
 from .vis import *
 
-__version__ = "0.4.1"
+__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 ..core.competitions import WTAC
-from ..core.components import Components, LabeledComponents
-from ..core.distances import euclidean_distance
-from ..core.initializers import LabelsInitializer
-from ..core.pooling import stratified_min_pooling
-from ..nn.wrappers import LambdaLayer
+from prototorch.core.competitions import WTAC
+from prototorch.core.components import (
+    AbstractComponents,
+    Components,
+    LabeledComponents,
+)
+from prototorch.core.distances import euclidean_distance
+from prototorch.core.initializers import (
+    LabelsInitializer,
+    ZerosCompInitializer,
+)
+from prototorch.core.pooling import stratified_min_pooling
+from prototorch.nn.wrappers import LambdaLayer
 
 
 class ProtoTorchBolt(pl.LightningModule):
@@ -30,7 +39,7 @@ class ProtoTorchBolt(pl.LightningModule):
         self.lr_scheduler_kwargs = kwargs.get("lr_scheduler_kwargs", dict())
 
     def configure_optimizers(self):
-        optimizer = self.optimizer(self.parameters(), lr=self.hparams.lr)
+        optimizer = self.optimizer(self.parameters(), lr=self.hparams["lr"])
         if self.lr_scheduler is not None:
             scheduler = self.lr_scheduler(optimizer,
                                           **self.lr_scheduler_kwargs)
@@ -43,7 +52,10 @@ class ProtoTorchBolt(pl.LightningModule):
             return optimizer
 
     def reconfigure_optimizers(self):
-        self.trainer.accelerator.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,14 +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.reconfigure_optimizers()
 
     def remove_prototypes(self, indices):
         self.proto_layer.remove_components(indices)
+        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)
@@ -91,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,
             )
 
@@ -106,20 +122,34 @@ class UnsupervisedPrototypeModel(PrototypeModel):
 
 
 class SupervisedPrototypeModel(PrototypeModel):
+    proto_layer: LabeledComponents
 
-    def __init__(self, hparams, **kwargs):
+    def __init__(self, hparams, skip_proto_layer=False, **kwargs):
         super().__init__(hparams, **kwargs)
 
         # Layers
+        distribution = hparams.get("distribution", None)
         prototypes_initializer = kwargs.get("prototypes_initializer", None)
         labels_initializer = kwargs.get("labels_initializer",
                                         LabelsInitializer())
+        if not skip_proto_layer:
+            # when subclasses do not need a customized prototype layer
             if prototypes_initializer is not None:
+                # when building a new model
                 self.proto_layer = LabeledComponents(
-                distribution=self.hparams.distribution,
+                    distribution=distribution,
                     components_initializer=prototypes_initializer,
                     labels_initializer=labels_initializer,
                 )
+                proto_shape = self.proto_layer.components.shape[1:]
+                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.competition_layer = WTAC()
 
     @property
@@ -139,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):
@@ -156,28 +186,38 @@ class SupervisedPrototypeModel(PrototypeModel):
 
     def log_acc(self, distances, targets, tag):
         preds = self.predict_from_distances(distances)
-        accuracy = torchmetrics.functional.accuracy(preds.int(), targets.int())
-        # `.int()` because FloatTensors are assumed to be class probabilities
+        accuracy = torchmetrics.functional.accuracy(
+            preds.int(),
+            targets.int(),
+            "multiclass",
+            num_classes=self.num_classes,
+        )
 
-        self.log(tag,
+        self.log(
+            tag,
             accuracy,
             on_step=False,
             on_epoch=True,
             prog_bar=True,
-                 logger=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."""
-    pass
 
 
 class NonGradientMixin(ProtoTorchMixin):
@@ -187,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__(
+        self,
         threshold=0.01,
         idle_epochs=10,
         prune_quota_per_epoch=-1,
         frequency=1,
         replace=False,
         prototypes_initializer=None,
-                 verbose=False):
+        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,26 +48,28 @@ class PruneLoserPrototypes(pl.Callback):
             prune_labels = prune_labels[:self.prune_quota_per_epoch]
 
         if len(to_prune) > 0:
-            if self.verbose:
-                print(f"\nPrototype win ratios: {ratios}")
-                print(f"Pruning prototypes at: {to_prune}")
-                print(f"Corresponding labels are: {prune_labels.tolist()}")
+            logging.debug(f"\nPrototype win ratios: {ratios}")
+            logging.debug(f"Pruning prototypes at: {to_prune}")
+            logging.debug(f"Corresponding labels are: {prune_labels.tolist()}")
+
             cur_num_protos = pl_module.num_prototypes
             pl_module.remove_prototypes(indices=to_prune)
+
             if self.replace:
                 labels, counts = torch.unique(prune_labels,
                                               sorted=True,
                                               return_counts=True)
                 distribution = dict(zip(labels.tolist(), counts.tolist()))
-                if self.verbose:
-                    print(f"Re-adding pruned prototypes...")
-                    print(f"distribution={distribution}")
+
+                logging.info(f"Re-adding pruned prototypes...")
+                logging.debug(f"distribution={distribution}")
+
                 pl_module.add_prototypes(
                     distribution=distribution,
                     components_initializer=self.prototypes_initializer)
             new_num_protos = pl_module.num_prototypes
-            if self.verbose:
-                print(f"`num_prototypes` changed from {cur_num_protos} "
+
+            logging.info(f"`num_prototypes` changed from {cur_num_protos} "
                          f"to {new_num_protos}.")
         return True
 
@@ -74,11 +81,11 @@ class PrototypeConvergence(pl.Callback):
         self.idle_epochs = idle_epochs  # epochs to wait
         self.verbose = verbose
 
-    def on_epoch_end(self, trainer, pl_module):
+    def on_train_epoch_end(self, trainer, pl_module):
         if (trainer.current_epoch + 1) < self.idle_epochs:
             return None
-        if self.verbose:
-            print("Stopping...")
+
+        logging.info("Stopping...")
         # TODO
         return True
 
@@ -96,12 +103,16 @@ class GNGCallback(pl.Callback):
         self.reduction = reduction
         self.freq = freq
 
-    def on_epoch_end(self, trainer: pl.Trainer, pl_module):
+    def on_train_epoch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: "GrowingNeuralGas",
+    ):
         if (trainer.current_epoch + 1) % self.freq == 0:
             # Get information
             errors = pl_module.errors
             topology: ConnectionTopology = pl_module.topology_layer
-            components: Components = pl_module.proto_layer.components
+            components = pl_module.proto_layer.components
 
             # Insertion point
             worst = torch.argmax(errors)
@@ -121,8 +132,9 @@ class GNGCallback(pl.Callback):
 
             # Add component
             pl_module.proto_layer.add_components(
-                None,
-                initializer=LiteralCompInitializer(new_component.unsqueeze(0)))
+                1,
+                initializer=LiteralCompInitializer(new_component.unsqueeze(0)),
+            )
 
             # Adjust Topology
             topology.add_prototype()
@@ -137,4 +149,4 @@ class GNGCallback(pl.Callback):
             pl_module.errors[
                 worst_neighbor] = errors[worst_neighbor] * self.reduction
 
-            trainer.accelerator.setup_optimizers(trainer)
+            trainer.strategy.setup_optimizers(trainer)

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
 
@@ -15,7 +15,7 @@ class CBC(SiameseGLVQ):
     """Classification-By-Components."""
 
     def __init__(self, hparams, **kwargs):
-        super().__init__(hparams, **kwargs)
+        super().__init__(hparams, skip_proto_layer=True, **kwargs)
 
         similarity_fn = kwargs.get("similarity_fn", euclidean_similarity)
         components_initializer = kwargs.get("components_initializer", None)
@@ -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):
-        y_pred, train_loss = self.shared_step(batch, batch_idx, optimizer_idx)
+    def training_step(self, batch, batch_idx):
+        y_pred, train_loss = self.shared_step(batch, batch_idx)
         preds = torch.argmax(y_pred, dim=1)
-        accuracy = torchmetrics.functional.accuracy(preds.int(),
-                                                    batch[1].int())
-        self.log("train_acc",
+        accuracy = torchmetrics.functional.accuracy(
+            preds.int(),
+            batch[1].int(),
+            "multiclass",
+            num_classes=self.num_classes,
+        )
+        self.log(
+            "train_acc",
             accuracy,
             on_step=False,
             on_epoch=True,
             prog_bar=True,
-                 logger=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 ..core.similarities import gaussian
+from prototorch.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 ..core.competitions import wtac
-from ..core.distances import (
+from prototorch.core.competitions import wtac
+from prototorch.core.distances import (
     lomega_distance,
     omega_distance,
     squared_euclidean_distance,
 )
-from ..core.initializers import EyeTransformInitializer
-from ..core.losses import (
+from prototorch.core.initializers import EyeLinearTransformInitializer
+from prototorch.core.losses import (
     GLVQLoss,
     lvq1_loss,
     lvq21_loss,
 )
-from ..core.transforms import LinearTransform
-from ..nn.wrappers import LambdaLayer, LossLayer
+from prototorch.core.transforms import LinearTransform
+from prototorch.nn.wrappers import LambdaLayer, LossLayer
+from torch.nn.parameter import Parameter
+
 from .abstract import ImagePrototypesMixin, SupervisedPrototypeModel
 from .extras import ltangent_distance, orthogonalization
 
@@ -34,17 +34,21 @@ class GLVQ(SupervisedPrototypeModel):
 
         # Loss
         self.loss = GLVQLoss(
-            margin=self.hparams.margin,
-            transfer_fn=self.hparams.transfer_fn,
-            beta=self.hparams.transfer_beta,
+            margin=self.hparams["margin"],
+            transfer_fn=self.hparams["transfer_fn"],
+            beta=self.hparams["transfer_beta"],
         )
 
+    # def on_save_checkpoint(self, checkpoint):
+    #     if "prototype_win_ratios" in checkpoint["state_dict"]:
+    #         del checkpoint["state_dict"]["prototype_win_ratios"]
+
     def initialize_prototype_win_ratios(self):
         self.register_buffer(
             "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):
@@ -62,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):
-        out, train_loss = self.shared_step(batch, batch_idx, optimizer_idx)
+    def training_step(self, batch, batch_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")
@@ -95,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.
@@ -119,33 +119,17 @@ 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)
-        self.backbone.requires_grad_(self.both_path_gradients)
+
+        bb_grad = any([el.requires_grad for el in self.backbone.parameters()])
+
+        self.backbone.requires_grad_(bb_grad and self.both_path_gradients)
         latent_protos = self.backbone(protos)
-        self.backbone.requires_grad_(True)
+        self.backbone.requires_grad_(bb_grad)
+
         distances = self.distance_layer(latent_x, latent_protos)
         return distances
 
@@ -191,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()
@@ -223,10 +211,10 @@ class SiameseGMLVQ(SiameseGLVQ):
 
         # Override the backbone
         omega_initializer = kwargs.get("omega_initializer",
-                                       EyeTransformInitializer())
+                                       EyeLinearTransformInitializer())
         self.backbone = LinearTransform(
-            self.hparams.input_dim,
-            self.hparams.output_dim,
+            self.hparams["input_dim"],
+            self.hparams["latent_dim"],
             initializer=omega_initializer,
         )
 
@@ -236,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()
 
@@ -249,18 +237,19 @@ 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)
 
         # Additional parameters
         omega_initializer = kwargs.get("omega_initializer",
-                                       EyeTransformInitializer())
-        omega = omega_initializer.generate(self.hparams.input_dim,
-                                           self.hparams.latent_dim)
+                                       EyeLinearTransformInitializer())
+        omega = omega_initializer.generate(self.hparams["input_dim"],
+                                           self.hparams["latent_dim"])
         self.register_parameter("_omega", Parameter(omega))
-        self.backbone = LambdaLayer(lambda x: x @ self._omega,
-                                    name="omega matrix")
 
     @property
     def omega_matrix(self):
@@ -291,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.latent_dim,
+            self.hparams["input_dim"],
+            self.hparams["latent_dim"],
             device=self.device,
         )
         self.register_parameter("_omega", Parameter(omega))
@@ -308,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,
-                                                  self.hparams.latent_dim)
-            omega = torch.repeat_interleave(subspace.unsqueeze(0),
+            subspace = omega_initializer.generate(
+                self.hparams["input_dim"],
+                self.hparams["latent_dim"],
+            )
+            omega = torch.repeat_interleave(
+                subspace.unsqueeze(0),
                 self.num_prototypes,
-                                            dim=0)
+                dim=0,
+            )
         else:
             omega = torch.rand(
                 self.num_prototypes,
-                self.hparams.input_dim,
-                self.hparams.latent_dim,
+                self.hparams["input_dim"],
+                self.hparams["latent_dim"],
                 device=self.device,
             )
 
         # Re-register `_omega` to override the one from the super class.
         self.register_parameter("_omega", Parameter(omega))
 
-    def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
+    def on_train_batch_end(self, outputs, batch, batch_idx):
         with torch.no_grad():
             self._omega.copy_(orthogonalization(self._omega))
 
@@ -381,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 ..core.components import LabeledComponents
-from ..core.initializers import (
+from prototorch.core.competitions import KNNC
+from prototorch.core.components import LabeledComponents
+from prototorch.core.initializers import (
     LiteralCompInitializer,
     LiteralLabelsInitializer,
 )
-from ..utils.utils import parse_data_arg
+from prototorch.utils.utils import parse_data_arg
+
 from .abstract import SupervisedPrototypeModel
 
 
@@ -16,7 +17,7 @@ class KNN(SupervisedPrototypeModel):
     """K-Nearest-Neighbors classification algorithm."""
 
     def __init__(self, hparams, **kwargs):
-        super().__init__(hparams, **kwargs)
+        super().__init__(hparams, skip_proto_layer=True, **kwargs)
 
         # Default hparams
         self.hparams.setdefault("k", 1)
@@ -28,18 +29,15 @@ class KNN(SupervisedPrototypeModel):
 
         # Layers
         self.proto_layer = LabeledComponents(
-            distribution=[],
+            distribution=len(data) * [1],
             components_initializer=LiteralCompInitializer(data),
             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,
-                             train_batch,
-                             batch_idx,
-                             dataloader_idx=None):
+    def on_train_batch_start(self, train_batch, batch_idx):
         warnings.warn("k-NN has no training, skipping!")
         return -1
 

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
-from ..nn.activations import get_activation
-from ..nn.wrappers import LambdaLayer
+import logging
+
+from prototorch.core.losses import _get_dp_dm
+from prototorch.nn.activations import get_activation
+from prototorch.nn.wrappers import LambdaLayer
+
 from .abstract import NonGradientMixin
 from .glvq import GLVQ
 
@@ -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}")
-        print(f"y={y}")
+        logging.debug(f"dis={dis}")
+        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):
-        self.verbose = verbose
+    def __init__(self, hparams, **kwargs):
         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:
-                        print(f"Updating prototype {i} to data {k}...")
+                    logging.debug(f"Updating prototype {i} to data {k}...")
                     self.proto_layer.load_state_dict({"_components": _protos},
                                                      strict=False)
                     break

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,14 +71,25 @@ 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."""
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+
+        # Default hparams
+        self.hparams.setdefault("variance", 1.0)
+        variance = self.hparams.get("variance")
+
+        self._conditional_distribution = GaussianPrior(variance)
         self.loss = LossLayer(nllr_loss)
-        self.conditional_distribution = GaussianPrior(self.hparams.variance)
 
 
 class RSLVQ(ProbabilisticLVQ):
@@ -76,8 +97,13 @@ class RSLVQ(ProbabilisticLVQ):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+
+        # Default hparams
+        self.hparams.setdefault("variance", 1.0)
+        variance = self.hparams.get("variance")
+
+        self._conditional_distribution = GaussianPrior(variance)
         self.loss = LossLayer(rslvq_loss)
-        self.conditional_distribution = GaussianPrior(self.hparams.variance)
 
 
 class PLVQ(ProbabilisticLVQ, SiameseGMLVQ):
@@ -88,12 +114,16 @@ class PLVQ(ProbabilisticLVQ, SiameseGMLVQ):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.conditional_distribution = RankScaledGaussianPrior(
-            self.hparams.lambd)
+
+        # Default hparams
+        self.hparams.setdefault("lambda", 1.0)
+        lam = self.hparams.get("lambda", 1.0)
+
+        self.conditional_distribution = RankScaledGaussianPrior(lam)
         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")
@@ -35,7 +35,7 @@ class KohonenSOM(NonGradientMixin, UnsupervisedPrototypeModel):
 
         # Additional parameters
         x, y = torch.arange(h), torch.arange(w)
-        grid = torch.stack(torch.meshgrid(x, y), dim=-1)
+        grid = torch.stack(torch.meshgrid(x, y, indexing="ij"), dim=-1)
         self.register_buffer("_grid", grid)
         self._sigma = self.hparams.sigma
         self._lr = self.hparams.lr
@@ -58,10 +58,12 @@ class KohonenSOM(NonGradientMixin, UnsupervisedPrototypeModel):
         diff = x.unsqueeze(dim=1) - protos
         delta = self._lr * self.hparams.alpha * nh.unsqueeze(-1) * diff
         updated_protos = protos + delta.sum(dim=0)
-        self.proto_layer.load_state_dict({"_components": updated_protos},
-                                         strict=False)
+        self.proto_layer.load_state_dict(
+            {"_components": updated_protos},
+            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)
 
@@ -88,13 +90,13 @@ class NeuralGas(UnsupervisedPrototypeModel):
         self.save_hyperparameters(hparams)
 
         # Default hparams
-        self.hparams.setdefault("agelimit", 10)
+        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.agelimit,
-            num_prototypes=self.hparams.num_prototypes,
+            agelimit=self.hparams["age_limit"],
+            num_prototypes=self.hparams["num_prototypes"],
         )
 
     def training_step(self, train_batch, batch_idx):
@@ -107,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)
@@ -122,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):
@@ -137,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)
@@ -145,6 +147,8 @@ class GrowingNeuralGas(NeuralGas):
 
     def configure_callbacks(self):
         return [
-            GNGCallback(reduction=self.hparams.insert_reduction,
-                        freq=self.hparams.insert_freq)
+            GNGCallback(
+                reduction=self.hparams["insert_reduction"],
+                freq=self.hparams["insert_freq"],
+            )
         ]

src/prototorch/models/vis.py

@@ -1,21 +1,28 @@
 """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.utils import mesh2d
-
 
 class Vis2DAbstract(pl.Callback):
 
     def __init__(self,
-                 data,
+                 data=None,
                  title="Prototype Visualization",
                  cmap="viridis",
+                 xlabel="Data dimension 1",
+                 ylabel="Data dimension 2",
+                 legend_labels=None,
                  border=0.1,
                  resolution=100,
                  flatten_data=True,
@@ -28,9 +35,15 @@ class Vis2DAbstract(pl.Callback):
                  block=False):
         super().__init__()
 
+        if data:
             if isinstance(data, Dataset):
+                if isinstance(data, Sized):
                     x, y = next(iter(DataLoader(data, batch_size=len(data))))
-        elif isinstance(data, torch.utils.data.DataLoader):
+                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:
@@ -44,8 +57,14 @@ class Vis2DAbstract(pl.Callback):
 
             self.x_train = x
             self.y_train = y
+        else:
+            self.x_train = None
+            self.y_train = None
 
         self.title = title
+        self.xlabel = xlabel
+        self.ylabel = ylabel
+        self.legend_labels = legend_labels
         self.fig = plt.figure(self.title)
         self.cmap = cmap
         self.border = border
@@ -64,14 +83,12 @@ class Vis2DAbstract(pl.Callback):
                 return False
         return True
 
-    def setup_ax(self, xlabel=None, ylabel=None):
+    def setup_ax(self):
         ax = self.fig.gca()
         ax.cla()
         ax.set_title(self.title)
-        if xlabel:
-            ax.set_xlabel("Data dimension 1")
-        if ylabel:
-            ax.set_ylabel("Data dimension 2")
+        ax.set_xlabel(self.xlabel)
+        ax.set_ylabel(self.ylabel)
         if self.axis_off:
             ax.axis("off")
         return ax
@@ -114,33 +131,39 @@ class Vis2DAbstract(pl.Callback):
             else:
                 plt.show(block=self.block)
 
+    def on_train_epoch_end(self, trainer, pl_module):
+        if not self.precheck(trainer):
+            return True
+        self.visualize(pl_module)
+        self.log_and_display(trainer, pl_module)
+
     def on_train_end(self, trainer, pl_module):
         plt.close()
 
+    def visualize(self, pl_module):
+        raise NotImplementedError
+
 
 class VisGLVQ2D(Vis2DAbstract):
 
-    def on_epoch_end(self, trainer, pl_module):
-        if not self.precheck(trainer):
-            return True
-
+    def visualize(self, pl_module):
         protos = pl_module.prototypes
         plabels = pl_module.prototype_labels
         x_train, y_train = self.x_train, self.y_train
-        ax = self.setup_ax(xlabel="Data dimension 1",
-                           ylabel="Data dimension 2")
-        self.plot_data(ax, x_train, y_train)
+        ax = self.setup_ax()
         self.plot_protos(ax, protos, plabels)
-        x = np.vstack((x_train, protos))
-        mesh_input, xx, yy = mesh2d(x, self.border, self.resolution)
+        if x_train is not None:
+            self.plot_data(ax, x_train, y_train)
+            mesh_input, xx, yy = mesh2d(np.vstack([x_train, protos]),
+                                        self.border, self.resolution)
+        else:
+            mesh_input, xx, yy = mesh2d(protos, self.border, self.resolution)
         _components = pl_module.proto_layer._components
         mesh_input = torch.from_numpy(mesh_input).type_as(_components)
         y_pred = pl_module.predict(mesh_input)
         y_pred = y_pred.cpu().reshape(xx.shape)
         ax.contourf(xx, yy, y_pred, cmap=self.cmap, alpha=0.35)
 
-        self.log_and_display(trainer, pl_module)
-
 
 class VisSiameseGLVQ2D(Vis2DAbstract):
 
@@ -148,10 +171,7 @@ class VisSiameseGLVQ2D(Vis2DAbstract):
         super().__init__(*args, **kwargs)
         self.map_protos = map_protos
 
-    def on_epoch_end(self, trainer, pl_module):
-        if not self.precheck(trainer):
-            return True
-
+    def visualize(self, pl_module):
         protos = pl_module.prototypes
         plabels = pl_module.prototype_labels
         x_train, y_train = self.x_train, self.y_train
@@ -178,8 +198,6 @@ class VisSiameseGLVQ2D(Vis2DAbstract):
         y_pred = y_pred.cpu().reshape(xx.shape)
         ax.contourf(xx, yy, y_pred, cmap=self.cmap, alpha=0.35)
 
-        self.log_and_display(trainer, pl_module)
-
 
 class VisGMLVQ2D(Vis2DAbstract):
 
@@ -187,10 +205,7 @@ class VisGMLVQ2D(Vis2DAbstract):
         super().__init__(*args, **kwargs)
         self.ev_proj = ev_proj
 
-    def on_epoch_end(self, trainer, pl_module):
-        if not self.precheck(trainer):
-            return True
-
+    def visualize(self, pl_module):
         protos = pl_module.prototypes
         plabels = pl_module.prototype_labels
         x_train, y_train = self.x_train, self.y_train
@@ -212,19 +227,13 @@ class VisGMLVQ2D(Vis2DAbstract):
         if self.show_protos:
             self.plot_protos(ax, protos, plabels)
 
-        self.log_and_display(trainer, pl_module)
-
 
 class VisCBC2D(Vis2DAbstract):
 
-    def on_epoch_end(self, trainer, pl_module):
-        if not self.precheck(trainer):
-            return True
-
+    def visualize(self, pl_module):
         x_train, y_train = self.x_train, self.y_train
         protos = pl_module.components
-        ax = self.setup_ax(xlabel="Data dimension 1",
-                           ylabel="Data dimension 2")
+        ax = self.setup_ax()
         self.plot_data(ax, x_train, y_train)
         self.plot_protos(ax, protos, "w")
         x = np.vstack((x_train, protos))
@@ -236,21 +245,15 @@ class VisCBC2D(Vis2DAbstract):
 
         ax.contourf(xx, yy, y_pred, cmap=self.cmap, alpha=0.35)
 
-        self.log_and_display(trainer, pl_module)
-
 
 class VisNG2D(Vis2DAbstract):
 
-    def on_epoch_end(self, trainer, pl_module):
-        if not self.precheck(trainer):
-            return True
-
+    def visualize(self, pl_module):
         x_train, y_train = self.x_train, self.y_train
         protos = pl_module.prototypes
         cmat = pl_module.topology_layer.cmat.cpu().numpy()
 
-        ax = self.setup_ax(xlabel="Data dimension 1",
-                           ylabel="Data dimension 2")
+        ax = self.setup_ax()
         self.plot_data(ax, x_train, y_train)
         self.plot_protos(ax, protos, "w")
 
@@ -264,7 +267,23 @@ class VisNG2D(Vis2DAbstract):
                         "k-",
                     )
 
-        self.log_and_display(trainer, pl_module)
+
+class VisSpectralProtos(Vis2DAbstract):
+
+    def visualize(self, pl_module):
+        protos = pl_module.prototypes
+        plabels = pl_module.prototype_labels
+        ax = self.setup_ax()
+        colors = get_colors(vmax=max(plabels), vmin=min(plabels))
+        for p, pl in zip(protos, plabels):
+            ax.plot(p, c=colors[int(pl)])
+        if self.legend_labels:
+            handles = get_legend_handles(
+                colors,
+                self.legend_labels,
+                marker="lines",
+            )
+            ax.legend(handles=handles)
 
 
 class VisImgComp(Vis2DAbstract):
@@ -284,9 +303,13 @@ 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):
+        if isinstance(pl_module.logger, TensorBoardLogger):
             tb = pl_module.logger.experiment
+
+            # Add embedding
             if self.add_embedding:
+                if self.x_train is not None and self.y_train is not None:
                     ind = np.random.choice(len(self.x_train),
                                            size=self.embedding_data,
                                            replace=False)
@@ -297,17 +320,28 @@ class VisImgComp(Vis2DAbstract):
                                      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)
+                    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
@@ -321,14 +355,9 @@ class VisImgComp(Vis2DAbstract):
             dataformats=self.dataformats,
         )
 
-    def on_epoch_end(self, trainer, pl_module):
-        if not self.precheck(trainer):
-            return True
-
+    def visualize(self, pl_module):
         if self.show:
             components = pl_module.components
             grid = torchvision.utils.make_grid(components,
                                                nrow=self.num_columns)
             plt.imshow(grid.permute((1, 2, 0)).cpu(), cmap=self.cmap)
-
-        self.log_and_display(trainer, pl_module)

tests/test_.py

@@ -1,15 +0,0 @@
-"""prototorch.models test suite."""
-
-import unittest
-
-
-class TestDummy(unittest.TestCase):
-
-    def setUp(self):
-        pass
-
-    def test_dummy(self):
-        pass
-
-    def tearDown(self):
-        pass

tests/test_models.py

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