Add more experimental changes

The code gets very messy very quickly as soon as serialization features are
needed.

examples/glvq_iris.py

@@ -25,6 +25,8 @@ class GLVQIris(GLVQ):
         parser.add_argument("--epochs", type=int, default=1)
         parser.add_argument("--lr", type=float, default=1e-1)
         parser.add_argument("--batch_size", type=int, default=150)
+        parser.add_argument("--input_dim", type=int, default=2)
+        parser.add_argument("--nclasses", type=int, default=3)
         parser.add_argument("--prototypes_per_class", type=int, default=3)
         parser.add_argument("--prototype_initializer",
                             type=str,
@@ -101,6 +103,7 @@ if __name__ == "__main__":
     # Setup trainer
     trainer = pl.Trainer.from_argparse_args(
         parser,
+        max_epochs=10,
         callbacks=[
             vis,  # comment this line out to disable the visualization
         ],
@@ -109,12 +112,7 @@ if __name__ == "__main__":
 
     # Initialize the model
     args = parser.parse_args()
-    model = GLVQIris(
-        args,
-        input_dim=x_train.shape[1],
-        nclasses=3,
-        data=[x_train, y_train],
-    )
+    model = GLVQIris(args, data=[x_train, y_train])
 
     # Model summary
     print(model)

prototorch/models/glvq.py

@@ -12,20 +12,14 @@ from prototorch.modules.prototypes import Prototypes1D
 
 class GLVQ(pl.LightningModule):
     """Generalized Learning Vector Quantization."""
-    def __init__(self, hparams, input_dim, nclasses, **kwargs):
+    def __init__(self, hparams, **kwargs):
         super().__init__()
-        self.lr = hparams.lr
-        self.hparams = hparams
-        # self.save_hyperparameters(
-        #     "lr",
-        #     "prototypes_per_class",
-        #     "prototype_initializer",
-        # )
+        self.save_hyperparameters(hparams)
         self.proto_layer = Prototypes1D(
-            input_dim=input_dim,
-            nclasses=nclasses,
-            prototypes_per_class=hparams.prototypes_per_class,
-            prototype_initializer=hparams.prototype_initializer,
+            input_dim=self.hparams.input_dim,
+            nclasses=self.hparams.nclasses,
+            prototypes_per_class=self.hparams.prototypes_per_class,
+            prototype_initializer=self.hparams.prototype_initializer,
             **kwargs)
         self.train_acc = torchmetrics.Accuracy()
 
@@ -38,22 +32,9 @@ class GLVQ(pl.LightningModule):
         return self.proto_layer.prototype_labels.detach().numpy()
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
         return optimizer
 
-    @staticmethod
-    def add_model_specific_args(parent_parser):
-        parser = argparse.ArgumentParser(parents=[parent_parser],
-                                         add_help=False)
-        parser.add_argument("--epochs", type=int, default=1)
-        parser.add_argument("--lr", type=float, default=1e-2)
-        parser.add_argument("--batch_size", type=int, default=32)
-        parser.add_argument("--prototypes_per_class", type=int, default=1)
-        parser.add_argument("--prototype_initializer",
-                            type=str,
-                            default="zeros")
-        return parser
-
     def forward(self, x):
         protos = self.proto_layer.prototypes
         dis = euclidean_distance(x, protos)