[QA] Add more pre commit checks

examples/glvq_iris.py

@@ -3,13 +3,14 @@
 import numpy as np
 import torch
 from matplotlib import pyplot as plt
+from sklearn.datasets import load_iris
+from sklearn.preprocessing import StandardScaler
+from torchinfo import summary
+
 from prototorch.components import LabeledComponents, StratifiedMeanInitializer
 from prototorch.functions.competitions import wtac
 from prototorch.functions.distances import euclidean_distance
 from prototorch.modules.losses import GLVQLoss
-from sklearn.datasets import load_iris
-from sklearn.preprocessing import StandardScaler
-from torchinfo import summary
 
 # Prepare and preprocess the data
 scaler = StandardScaler()

examples/gmlvq_tecator.py

@@ -2,12 +2,13 @@
 
 import matplotlib.pyplot as plt
 import torch
+from torch.utils.data import DataLoader
+
 from prototorch.components import LabeledComponents, StratifiedMeanInitializer
 from prototorch.datasets.tecator import Tecator
 from prototorch.functions.distances import sed
 from prototorch.modules.losses import GLVQLoss
 from prototorch.utils.colors import get_legend_handles
-from torch.utils.data import DataLoader
 
 # Prepare the dataset and dataloader
 train_data = Tecator(root="./artifacts", train=True)

examples/gtlvq_mnist.py

@@ -12,10 +12,11 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torchvision
+from torchvision import transforms
+
 from prototorch.functions.helper import calculate_prototype_accuracy
 from prototorch.modules.losses import GLVQLoss
 from prototorch.modules.models import GTLVQ
-from torchvision import transforms
 
 # Parameters and options
 num_epochs = 50

examples/lgmlvq_iris.py

@@ -3,13 +3,14 @@
 import numpy as np
 import torch
 from matplotlib import pyplot as plt
-from prototorch.components import LabeledComponents, StratifiedMeanInitializer
-from prototorch.functions.competitions import stratified_min
-from prototorch.functions.distances import lomega_distance
-from prototorch.modules.losses import GLVQLoss
 from sklearn.datasets import load_iris
 from sklearn.metrics import accuracy_score
 
+from prototorch.components import LabeledComponents, StratifiedMeanInitializer
+from prototorch.functions.distances import lomega_distance
+from prototorch.functions.pooling import stratified_min_pooling
+from prototorch.modules.losses import GLVQLoss
+
 # Prepare training data
 x_train, y_train = load_iris(True)
 x_train = x_train[:, [0, 2]]
@@ -55,7 +56,8 @@ for epoch in range(100):
     # Compute loss
     dis, plabels = model(x_in)
     loss = criterion([dis, plabels], y_in)
-    y_pred = np.argmin(stratified_min(dis, plabels).detach().numpy(), axis=1)
+    y_pred = np.argmin(stratified_min_pooling(dis, plabels).detach().numpy(),
+                       axis=1)
     acc = accuracy_score(y_train, y_pred)
     log_string = f"Epoch: {epoch + 1:03d} Loss: {loss.item():05.02f} "
     log_string += f"Acc: {acc * 100:05.02f}%"
@@ -96,7 +98,8 @@ for epoch in range(100):
     mesh_input = np.c_[xx.ravel(), yy.ravel()]
 
     d, plabels = model(torch.Tensor(mesh_input))
-    y_pred = np.argmin(stratified_min(d, plabels).detach().numpy(), axis=1)
+    y_pred = np.argmin(stratified_min_pooling(d, plabels).detach().numpy(),
+                       axis=1)
     y_pred = y_pred.reshape(xx.shape)
 
     # Plot voronoi regions