accelerator examples

Author: pereverges

examples/mnist_hugging_face.py

@@ -0,0 +1,85 @@
+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
+from accelerate import Accelerator
+
+from torchhd import functional
+from torchhd import embeddings
+
+accelerator = Accelerator()
+device = accelerator.device
+print("Using {} device".format(device))
+
+DIMENSIONS = 10000
+IMG_SIZE = 28
+NUM_LEVELS = 1000
+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.position = embeddings.Random(size * size, DIMENSIONS)
+        self.value = embeddings.Level(NUM_LEVELS, 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 = functional.bind(self.position.weight, self.value(x))
+        sample_hv = functional.multiset(sample_hv)
+        return functional.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.to(device)
+
+model, train_ld, test_ld = accelerator.prepare(
+    model, train_ld, test_ld
+)
+
+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}%")

examples/mnist_torch_lightning.py

@@ -0,0 +1,90 @@
+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 pytorch_lightning as pl
+from torchhd import functional
+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
+NUM_LEVELS = 1000
+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(pl.LightningModule):
+    def __init__(self, num_classes, size):
+        super(Model, self).__init__()
+
+        self.flatten = torch.nn.Flatten()
+
+        self.position = embeddings.Random(size * size, DIMENSIONS)
+        self.value = embeddings.Level(NUM_LEVELS, 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 = functional.bind(self.position.weight, self.value(x))
+        sample_hv = functional.multiset(sample_hv)
+        return functional.hard_quantize(sample_hv)
+
+    def forward(self, x):
+        enc = self.encode(x)
+        logit = self.classify(enc)
+        return logit
+
+    def training_step(self, batch):
+        return
+
+    def configure_optimizers(self):
+        return
+
+
+model = Model(len(train_ds.classes), IMG_SIZE)
+trainer = pl.Trainer(
+    accelerator="cpu",
+    devices=1,
+)
+trainer.fit(model, train_ld, test_ld)
+
+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}%")