Update example scripts

examples/glvq_spiral.py

@@ -2,11 +2,10 @@
 
 import argparse
 
+import prototorch as pt
 import pytorch_lightning as pl
 import torch
 
-import prototorch as pt
-
 if __name__ == "__main__":
     # Command-line arguments
     parser = argparse.ArgumentParser()
@@ -14,35 +13,63 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     # Dataset
-    train_ds = pt.datasets.Spiral(num_samples=600, noise=0.6)
+    train_ds = pt.datasets.Spiral(num_samples=500, noise=0.5)
 
     # Dataloaders
-    train_loader = torch.utils.data.DataLoader(train_ds,
-                                               num_workers=0,
-                                               batch_size=256)
+    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=256)
 
     # Hyperparameters
     num_classes = 2
-    prototypes_per_class = 20
+    prototypes_per_class = 10
     hparams = dict(
         distribution=(num_classes, prototypes_per_class),
-        transfer_function="sigmoid_beta",
+        transfer_function="swish_beta",
         transfer_beta=10.0,
-        lr=0.01,
+        # lr=0.1,
+        proto_lr=0.1,
+        bb_lr=0.1,
+        input_dim=2,
+        latent_dim=2,
     )
 
     # Initialize the model
-    model = pt.models.GLVQ(hparams,
-                           prototype_initializer=pt.components.SSI(train_ds,
-                                                                   noise=1e-1))
+    model = pt.models.GMLVQ(
+        hparams,
+        optimizer=torch.optim.Adam,
+        prototype_initializer=pt.components.SSI(train_ds, noise=1e-2),
+    )
 
     # Callbacks
-    vis = pt.models.VisGLVQ2D(train_ds, show_last_only=True, block=True)
+    vis = pt.models.VisGLVQ2D(
+        train_ds,
+        show_last_only=False,
+        block=False,
+    )
+    pruning = pt.models.PruneLoserPrototypes(
+        threshold=0.02,
+        idle_epochs=10,
+        prune_quota_per_epoch=5,
+        frequency=2,
+        replace=True,
+        initializer=pt.components.SSI(train_ds, noise=1e-2),
+        verbose=True,
+    )
+    es = pl.callbacks.EarlyStopping(
+        monitor="train_loss",
+        min_delta=1.0,
+        patience=5,
+        mode="min",
+        check_on_train_epoch_end=True,
+    )
 
     # Setup trainer
     trainer = pl.Trainer.from_argparse_args(
         args,
-        callbacks=[vis],
+        callbacks=[
+            vis,
+            # es,
+            pruning,
+        ],
         terminate_on_nan=True,
     )