Fix the nonlinear bug and add mnist example (#110)

* fix the nonlinear bug and add mnist example

* remove unused variable and add one comment

Author: Orienfish

examples/mnist_nonlinear.py

@@ -0,0 +1,77 @@
+# This is an example of using nonlinear encoding on the MNIST dataset
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+from torchvision.datasets import MNIST
+
+# Note: this example requires the torchmetrics library: https://torchmetrics.readthedocs.io
+import torchmetrics
+from tqdm import tqdm
+
+import torchhd
+from torchhd import embeddings
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using {} device".format(device))
+
+DIMENSIONS = 10000
+IMG_SIZE = 28
+BATCH_SIZE = 1  # for GPUs with enough memory we can process multiple images at ones
+
+transform = torchvision.transforms.ToTensor()
+
+train_ds = MNIST("../data", train=True, transform=transform, download=True)
+train_ld = torch.utils.data.DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
+
+test_ds = MNIST("../data", train=False, transform=transform, download=True)
+test_ld = torch.utils.data.DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False)
+
+class Model(nn.Module):
+    def __init__(self, num_classes, size):
+        super(Model, self).__init__()
+
+        self.flatten = torch.nn.Flatten()
+
+        self.nonlinear_projection = embeddings.Sinusoid(size * size, DIMENSIONS)
+
+        self.classify = nn.Linear(DIMENSIONS, num_classes, bias=False)
+        self.classify.weight.data.fill_(0.0)
+
+    def encode(self, x):
+        x = self.flatten(x)
+        sample_hv = self.nonlinear_projection(x)
+        return torchhd.hard_quantize(sample_hv)
+
+    def forward(self, x):
+        enc = self.encode(x)
+        logit = self.classify(enc)
+        return logit
+
+
+model = Model(len(train_ds.classes), IMG_SIZE)
+model = model.to(device)
+
+with torch.no_grad():
+    for samples, labels in tqdm(train_ld, desc="Training"):
+        samples = samples.to(device)
+        labels = labels.to(device)
+
+        samples_hv = model.encode(samples)
+        model.classify.weight[labels] += samples_hv
+
+    model.classify.weight[:] = F.normalize(model.classify.weight)
+
+accuracy = torchmetrics.Accuracy()
+
+
+with torch.no_grad():
+    for samples, labels in tqdm(test_ld, desc="Testing"):
+        samples = samples.to(device)
+
+        outputs = model(samples)
+        predictions = torch.argmax(outputs, dim=-1)
+        accuracy.update(predictions.cpu(), labels)
+
+print(f"Testing accuracy of {(accuracy.compute().item() * 100):.3f}%")

torchhd/embeddings.py

@@ -409,6 +409,7 @@ def __init__(
 
     def reset_parameters(self) -> None:
         nn.init.normal_(self.weight, 0, 1)
+        self.weight.data.copy_(F.normalize(self.weight.data))
         nn.init.uniform_(self.bias, 0, 2 * math.pi)
 
     def forward(self, input: torch.Tensor) -> torch.Tensor: