diff --git a/Norse/CIFAR-10/CIFAR10_conv_bench.py b/Norse/CIFAR-10/CIFAR10_conv_bench.py
new file mode 100644
index 0000000000000000000000000000000000000000..9832ec31d06d0a53543d8fd6869e89c68496cdef
--- /dev/null
+++ b/Norse/CIFAR-10/CIFAR10_conv_bench.py
@@ -0,0 +1,219 @@
+import sys
+sys.path.append('../../')
+
+import torch
+import numpy as np
+from compression import ProgressiveCompression
+from torch.utils.data import SubsetRandomSampler
+from norse.torch import LIFCell, LICell
+from norse.torch.module.leaky_integrator import LILinearCell
+from norse.torch import LIFParameters
+from norse.torch.module import encode, SequentialState
+from datetime import datetime
+import torchvision
+import os
+import pickle
+import random
+import itertools
+
+# Reproducibility
+torch.manual_seed(0)
+random.seed(0)
+np.random.seed(0)
+
+MAXTH = [0.3,0.4,0.5,0.6,0.7] #
+ALPHA = [0.005] # [0.002,0.004,0.006,0.008,0.01] #
+REINFORCEMENT = [True] # [False, True]
+apply_compression = False
+
+for maxTh, Alpha, reinforcement in np.array(list(itertools.product(MAXTH, ALPHA, REINFORCEMENT))):
+ try:
+ os.mkdir("CIFAR10_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement))
+ except OSError as error:
+ print(error)
+
+ for i in range(2):
+ before = datetime.now()
+
+ file = open("CIFAR10_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/CIFAR10_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before), 'w+')
+
+ transform = torchvision.transforms.Compose(
+ [
+ torchvision.transforms.ToTensor(),
+ torchvision.transforms.Normalize((0.1307,), (0.3081,)),
+ ]
+ )
+
+ train_data = torchvision.datasets.CIFAR10(
+ root=".",
+ train=True,
+ download=True,
+ transform=transform,
+ )
+
+ # reduce this number if you run out of GPU memory
+ BATCH_SIZE = 32
+
+ train_loader = torch.utils.data.DataLoader(
+ train_data, batch_size=BATCH_SIZE, shuffle=True #, sampler=SubsetRandomSampler(list(range(len(train_data)))[0:1000]) #
+ )
+
+ test_loader = torch.utils.data.DataLoader(
+ torchvision.datasets.CIFAR10(
+ root=".",
+ train=False,
+ transform=transform
+ ),
+ batch_size=BATCH_SIZE
+ )
+
+ class Model(torch.nn.Module):
+ def __init__(self, encoder, snn, decoder):
+ super(Model, self).__init__()
+ self.encoder = encoder
+ self.snn = snn
+ self.decoder = decoder
+
+ def forward(self, x):
+ x = self.encoder(x)
+ x = self.snn(x)
+ log_p_y = self.decoder(x)
+ return log_p_y
+
+ class ConvNet(torch.nn.Module):
+ def __init__(self, num_channels=3, feature_size=32, method="super", alpha=100):
+ super(ConvNet, self).__init__()
+
+ self.features = int(((feature_size - 4) / 2 - 4) / 2)
+ self.conv1_out_channels = 32
+ self.conv2_out_channels = 128
+ self.fc1_out_channels = 1024
+ self.out_channels = 10
+ self.conv1 = torch.nn.Conv2d(num_channels, self.conv1_out_channels, 5, 1, bias=False)
+ self.conv2 = torch.nn.Conv2d(self.conv1_out_channels, self.conv2_out_channels, 5, 1, bias=False)
+ self.fc1 = torch.nn.Linear(self.features**2 * self.conv2_out_channels, self.fc1_out_channels, bias=False)
+ self.lif0 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.25))
+ self.lif1 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.25))
+ self.lif2 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.25))
+ self.out = LILinearCell(self.fc1_out_channels, self.out_channels)
+
+ #LILinearCell(self.fc1_out_channels, self.out_channels)
+
+ def forward(self, x):
+ seq_length = x.shape[0]
+ batch_size = x.shape[1]
+
+ # specify the initial states
+ s0 = s1 = s2 = s3 = s4 = so = None
+
+ voltages = torch.zeros(
+ seq_length, batch_size, self.out_channels, device=x.device, dtype=x.dtype
+ )
+
+ for ts in range(seq_length):
+ z = self.conv1(x[ts, :])
+ z, s0 = self.lif0(z, s0)
+ z = torch.nn.functional.max_pool2d(z, 2, 2)
+ z = self.out_channels * self.conv2(z)
+ z, s1 = self.lif1(z, s1)
+ z = torch.nn.functional.max_pool2d(z, 2, 2)
+ z = z.view(-1, self.features**2 * self.conv2_out_channels)
+ z = self.fc1(z)
+ z, s2 = self.lif2(z, s2)
+ v, so = self.out(torch.nn.functional.relu(z), so)
+ voltages[ts, :, :] = v
+
+ return voltages
+
+ def train(model, device, train_loader, optimizer, epoch, max_epochs):
+ model.train()
+ losses = []
+
+ for (data, target) in train_loader:
+ data, target = data.to(device), target.to(device)
+ optimizer.zero_grad()
+ output = model(data)
+ loss = torch.nn.functional.nll_loss(output, target)
+ loss.backward()
+ optimizer.step()
+ losses.append(loss.item())
+
+ mean_loss = np.mean(losses)
+ return losses, mean_loss
+
+ def test(model, device, test_loader, epoch):
+ model.eval()
+ test_loss = 0
+ correct = 0
+ with torch.no_grad():
+ for data, target in test_loader:
+ data, target = data.to(device), target.to(device)
+ output = model(data)
+ test_loss += torch.nn.functional.nll_loss(
+ output, target, reduction="sum"
+ ).item() # sum up batch loss
+ pred = output.argmax(
+ dim=1, keepdim=True
+ ) # get the index of the max log-probability
+ correct += pred.eq(target.view_as(pred)).sum().item()
+
+ test_loss /= len(test_loader.dataset)
+
+ accuracy = 100.0 * correct / len(test_loader.dataset)
+
+ return test_loss, accuracy
+
+ def decode(x):
+ x, _ = torch.max(x, 0)
+ log_p_y = torch.nn.functional.log_softmax(x, dim=1)
+ return log_p_y
+
+ torch.autograd.set_detect_anomaly(True)
+
+ T = 35
+ LR = 0.001
+ EPOCHS = 100 # Increase this for improved accuracy
+
+ if torch.cuda.is_available():
+ DEVICE = torch.device(sys.argv[1])
+ else:
+ DEVICE = torch.device("cpu")
+
+ model = Model(
+ encoder=encode.SpikeLatencyLIFEncoder(T), snn=ConvNet(alpha=80), decoder=decode).to(DEVICE)
+
+ optimizer = torch.optim.Adam(model.parameters(), lr=LR)
+
+ # compression
+ if (apply_compression):
+ progressive_compression = ProgressiveCompression(NorseModel=model, maxThreshold=maxTh, alphaP=Alpha, alphaN=-Alpha, to_file=True, apply_reinforcement=reinforcement, file= file)
+
+ training_losses = []
+ mean_losses = []
+ test_losses = []
+ accuracies = []
+
+ for epoch in range(EPOCHS):
+ print(f"Epoch {epoch}")
+ training_loss, mean_loss = train(
+ model, DEVICE, train_loader, optimizer, epoch, max_epochs=EPOCHS
+ )
+ test_loss, accuracy = test(model, DEVICE, test_loader, epoch)
+ training_losses += training_loss
+ mean_losses.append(mean_loss)
+ test_losses.append(test_loss)
+ accuracies.append(accuracy)
+ if (apply_compression):
+ progressive_compression.apply()
+
+ print(f"final accuracy: {accuracies[-1]}")
+ file.write("final accuracy:"+str(accuracies[-1])+"\n")
+ file.write("time:"+str(datetime.now() - before)+"\n")
+
+
+ with open("CIFAR10_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/CIFAR10_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before)+".pkl",'wb') as f:
+ torch.save(model,"CIFAR10_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/CIFAR10_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before)+".norse")
+ if (apply_compression):
+ pickle.dump([mean_losses,test_losses,accuracies,progressive_compression.weights,progressive_compression.compressions,progressive_compression.thresholds_p,progressive_compression.thresholds_n], f)
+ else:
+ pickle.dump([mean_losses,test_losses,accuracies], f)
diff --git a/Norse/Caltech Face-Motor/CaltechFace_Motor_conv.py b/Norse/Caltech Face-Motor/CaltechFace_Motor_conv.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ddf1a39fcd72d87a1af4d6a02a3fdc8c3e4c67f
--- /dev/null
+++ b/Norse/Caltech Face-Motor/CaltechFace_Motor_conv.py
@@ -0,0 +1,214 @@
+import sys
+sys.path.append('../../')
+
+import torch
+import numpy as np
+from compression import ProgressiveCompression
+from norse.torch import LIFCell
+from norse.torch.module.leaky_integrator import LILinearCell
+from norse.torch import LIFParameters
+from norse.torch.module import encode
+from tqdm import tqdm, trange
+from datetime import datetime
+import torchvision
+import os
+import pickle
+import random
+import itertools
+
+# Reproducibility
+torch.manual_seed(0)
+random.seed(0)
+np.random.seed(0)
+
+MAXTH = [0.3,0.4,0.5,0.6,0.7] #
+ALPHA = [0.005] # [0.002,0.004,0.006,0.008,0.01] #
+REINFORCEMENT = [True] # [False, True]
+apply_compression = True
+
+for maxTh, Alpha, reinforcement in np.array(list(itertools.product(MAXTH, ALPHA, REINFORCEMENT))):
+ try:
+ os.mkdir("FACEMOTOR_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement))
+ except OSError as error:
+ print(error)
+
+ for i in range(2):
+ before = datetime.now()
+
+ file = open("FACEMOTOR_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/FACEMOTOR_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before), 'w+')
+
+ transform = torchvision.transforms.Compose(
+ [
+ torchvision.transforms.ToTensor(),
+ torchvision.transforms.Grayscale(),
+ torchvision.transforms.Resize((250,160)),
+ torchvision.transforms.Normalize((0.1307,), (0.3081,)),
+ ]
+ )
+
+ train_data = torchvision.datasets.ImageFolder(
+ root="/home/hammouda/Desktop/Work/falez-csnn-simulator/Datasets/FaceMotor/TrainingSet/",
+ transform=transform
+ )
+
+ # reduce this number if you run out of GPU memory
+ BATCH_SIZE = 5
+
+ train_loader = torch.utils.data.DataLoader(
+ train_data, batch_size=BATCH_SIZE, shuffle=True
+ )
+
+ test_loader = torch.utils.data.DataLoader(
+ torchvision.datasets.ImageFolder(
+ root="/home/hammouda/Desktop/Work/falez-csnn-simulator/Datasets/FaceMotor/TestingSet/",
+ transform=transform
+ ),
+ batch_size=BATCH_SIZE,
+ )
+
+ class Model(torch.nn.Module):
+ def __init__(self, encoder, snn, decoder):
+ super(Model, self).__init__()
+ self.encoder = encoder
+ self.snn = snn
+ self.decoder = decoder
+
+ def forward(self, x):
+ x = self.encoder(x)
+ x = self.snn(x)
+ log_p_y = self.decoder(x)
+ return log_p_y
+
+ class ConvNet(torch.nn.Module):
+ def __init__(self, num_channels=1, feature_size=250, method="super", alpha=100):
+ super(ConvNet, self).__init__()
+
+ self.features = 54
+ self.conv1_out_channels = 32
+ self.conv2_out_channels = 64
+ self.fc1_out_channels = 128
+ self.out_channels = 10
+ self.conv1 = torch.nn.Conv2d(num_channels, self.conv1_out_channels, 5, 1, 3, bias=False)
+ self.conv2 = torch.nn.Conv2d(self.conv1_out_channels, self.conv2_out_channels, 17, 1, 9, bias=False)
+ self.fc1 = torch.nn.Linear(self.features * self.conv2_out_channels, self.fc1_out_channels, bias=False)
+ self.lif0 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.25))
+ self.lif1 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.25))
+ self.lif2 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.25))
+ self.out = LILinearCell(self.fc1_out_channels, self.out_channels)
+
+ def forward(self, x):
+ seq_length = x.shape[0]
+ batch_size = x.shape[1]
+
+ # specify the initial states
+ s0 = s1 = s2 = so = None
+
+ voltages = torch.zeros(
+ seq_length, batch_size, self.out_channels, device=x.device, dtype=x.dtype
+ )
+
+ for ts in range(seq_length):
+ z = self.conv1(x[ts, :])
+ z, s0 = self.lif0(z, s0)
+ z = torch.nn.functional.max_pool2d(z, 7, 6, 3)
+ z = self.out_channels * self.conv2(z)
+ z, s1 = self.lif1(z, s1)
+ z = torch.nn.functional.max_pool2d(z, 5, 5, 2)
+ z = z.view(batch_size,-1)
+ z = self.fc1(z)
+ z, s2 = self.lif2(z, s2)
+ v, so = self.out(torch.nn.functional.relu(z), so)
+ voltages[ts, :, :] = v
+
+ return voltages
+
+ def train(model, device, train_loader, optimizer, epoch, max_epochs):
+ model.train()
+ losses = []
+
+ for (data, target) in tqdm(train_loader, leave=False):
+ data, target = data.to(device), target.to(device)
+ optimizer.zero_grad()
+ output = model(data)
+ loss = torch.nn.functional.nll_loss(output, target)
+ loss.backward()
+ optimizer.step()
+ losses.append(loss.item())
+
+ mean_loss = np.mean(losses)
+ return losses, mean_loss
+
+ def test(model, device, test_loader, epoch):
+ model.eval()
+ test_loss = 0
+ correct = 0
+ with torch.no_grad():
+ for data, target in test_loader:
+ data, target = data.to(device), target.to(device)
+ output = model(data)
+ test_loss += torch.nn.functional.nll_loss(
+ output, target, reduction="sum"
+ ).item() # sum up batch loss
+ pred = output.argmax(
+ dim=1, keepdim=True
+ ) # get the index of the max log-probability
+ correct += pred.eq(target.view_as(pred)).sum().item()
+
+ test_loss /= len(test_loader.dataset)
+
+ accuracy = 100.0 * correct / len(test_loader.dataset)
+
+ return test_loss, accuracy
+
+ def decode(x):
+ x, _ = torch.max(x, 0)
+ log_p_y = torch.nn.functional.log_softmax(x, dim=1)
+ return log_p_y
+
+ T = 35
+ LR = 1e-4
+ EPOCHS = 15 # Increase this for improved accuracy
+
+ if torch.cuda.is_available():
+ DEVICE = torch.device("cuda")
+ else:
+ DEVICE = torch.device("cpu")
+
+ model = Model(
+ encoder=encode.SpikeLatencyLIFEncoder(T), snn=ConvNet(alpha=80), decoder=decode).to(DEVICE)
+
+ optimizer = torch.optim.Adam(model.parameters(), lr=LR)
+
+ # compression
+ if (apply_compression):
+ progressive_compression = ProgressiveCompression(NorseModel=model, maxThreshold=maxTh, alphaP=Alpha, alphaN=-Alpha, to_file=True, apply_reinforcement=reinforcement, file= file)
+
+ training_losses = []
+ mean_losses = []
+ test_losses = []
+ accuracies = []
+
+ for epoch in trange(EPOCHS):
+ print(f"Epoch {epoch}")
+ training_loss, mean_loss = train(
+ model, DEVICE, train_loader, optimizer, epoch, max_epochs=EPOCHS
+ )
+ test_loss, accuracy = test(model, DEVICE, test_loader, epoch)
+ training_losses += training_loss
+ mean_losses.append(mean_loss)
+ test_losses.append(test_loss)
+ accuracies.append(accuracy)
+ if (apply_compression):
+ progressive_compression.apply()
+
+ print(f"final accuracy: {accuracies[-1]}")
+ file.write("final accuracy:"+str(accuracies[-1])+"\n")
+ file.write("time:"+str(datetime.now() - before)+"\n")
+
+
+ with open("FACEMOTOR_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/FACEMOTOR_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before)+".pkl",'wb') as f:
+ torch.save(model,"FACEMOTOR_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/FACEMOTOR_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before)+".norse")
+ if (apply_compression):
+ pickle.dump([mean_losses,test_losses,accuracies,progressive_compression.weights,progressive_compression.compressions,progressive_compression.thresholds_p,progressive_compression.thresholds_n], f)
+ else:
+ pickle.dump([mean_losses,test_losses,accuracies], f)
diff --git a/Norse/FMNIST/FMNIST_conv_bench.py b/Norse/FMNIST/FMNIST_conv_bench.py
new file mode 100644
index 0000000000000000000000000000000000000000..c3fb080ef867e67cd7fb0b0a0435e8b50ee57288
--- /dev/null
+++ b/Norse/FMNIST/FMNIST_conv_bench.py
@@ -0,0 +1,215 @@
+import sys
+sys.path.append('../../')
+
+import torch
+import numpy as np
+from compression import ProgressiveCompression
+from norse.torch import LIFCell
+from norse.torch.module.leaky_integrator import LILinearCell
+from norse.torch import LIFParameters
+from norse.torch.module import encode
+#from tqdm import tqdm, trange
+from datetime import datetime
+import torchvision
+import os
+import pickle
+import random
+import itertools
+
+# Reproducibility
+torch.manual_seed(0)
+random.seed(0)
+np.random.seed(0)
+
+MAXTH = [0.6,0.7] # 0.3,0.5,
+ALPHA = [0.005] # [0.002,0.004,0.006,0.008,0.01] #
+REINFORCEMENT = [True] # [False, True]
+
+for maxTh, Alpha, reinforcement in np.array(list(itertools.product(MAXTH, ALPHA, REINFORCEMENT))):
+ try:
+ os.mkdir("FashionMNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement))
+ except OSError as error:
+ print(error)
+
+ for i in range(2):
+ before = datetime.now()
+
+ file = open("FashionMNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/FashionMNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before), 'w+')
+
+ apply_compression = True
+
+ transform = torchvision.transforms.Compose(
+ [
+ torchvision.transforms.ToTensor(),
+ torchvision.transforms.Normalize((0.1307,), (0.3081,)),
+ ]
+ )
+
+ train_data = torchvision.datasets.FashionMNIST(
+ root=".",
+ train=True,
+ download=True,
+ transform=transform,
+ )
+
+ # reduce this number if you run out of GPU memory
+ BATCH_SIZE = 128
+
+ train_loader = torch.utils.data.DataLoader(
+ train_data, batch_size=BATCH_SIZE, shuffle=True
+ )
+
+ test_loader = torch.utils.data.DataLoader(
+ torchvision.datasets.FashionMNIST(
+ root=".",
+ train=False,
+ transform=transform,
+ ),
+ batch_size=BATCH_SIZE,
+ )
+
+ class Model(torch.nn.Module):
+ def __init__(self, encoder, snn, decoder):
+ super(Model, self).__init__()
+ self.encoder = encoder
+ self.snn = snn
+ self.decoder = decoder
+
+ def forward(self, x):
+ x = self.encoder(x)
+ x = self.snn(x)
+ log_p_y = self.decoder(x)
+ return log_p_y
+
+ class ConvNet(torch.nn.Module):
+ def __init__(self, num_channels=1, feature_size=28, method="super", alpha=100):
+ super(ConvNet, self).__init__()
+
+ self.features = int(((feature_size - 4) / 2 - 4) / 2)
+ self.conv1_out_channels = 32
+ self.conv2_out_channels = 128
+ self.fc1_out_channels = 1024
+ self.out_channels = 10
+ self.conv1 = torch.nn.Conv2d(num_channels, self.conv1_out_channels, 5, 1, bias=False)
+ self.conv2 = torch.nn.Conv2d(self.conv1_out_channels, self.conv2_out_channels, 5, 1, bias=False)
+ self.fc1 = torch.nn.Linear(self.features * self.features * self.conv2_out_channels, self.fc1_out_channels, bias=False)
+ self.lif0 = LIFCell(p=LIFParameters(method=method, alpha=alpha,v_th=0.25))
+ self.lif1 = LIFCell(p=LIFParameters(method=method, alpha=alpha,v_th=0.25))
+ self.lif2 = LIFCell(p=LIFParameters(method=method, alpha=alpha,v_th=0.25))
+ self.out = LILinearCell(self.fc1_out_channels, self.out_channels)
+
+ def forward(self, x):
+ seq_length = x.shape[0]
+ batch_size = x.shape[1]
+
+ # specify the initial states
+ s0 = s1 = s2 = so = None
+
+ voltages = torch.zeros(
+ seq_length, batch_size, self.out_channels, device=x.device, dtype=x.dtype
+ )
+
+ for ts in range(seq_length):
+ z = self.conv1(x[ts, :])
+ z, s0 = self.lif0(z, s0)
+ z = torch.nn.functional.max_pool2d(z, 2, 2)
+ z = self.out_channels * self.conv2(z)
+ z, s1 = self.lif1(z, s1)
+ z = torch.nn.functional.max_pool2d(z, 2, 2)
+ z = z.view(-1, 4**2 * self.conv2_out_channels)
+ z = self.fc1(z)
+ z, s2 = self.lif2(z, s2)
+ v, so = self.out(torch.nn.functional.relu(z), so)
+ voltages[ts, :, :] = v
+ return voltages
+
+ def train(model, device, train_loader, optimizer, epoch, max_epochs):
+ model.train()
+ losses = []
+
+ for (data, target) in train_loader: #tqdm(train_loader, leave=False):
+ data, target = data.to(device), target.to(device)
+ optimizer.zero_grad()
+ output = model(data)
+ loss = torch.nn.functional.nll_loss(output, target)
+ loss.backward()
+ optimizer.step()
+ losses.append(loss.item())
+
+ mean_loss = np.mean(losses)
+ return losses, mean_loss
+
+ def test(model, device, test_loader, epoch):
+ model.eval()
+ test_loss = 0
+ correct = 0
+ with torch.no_grad():
+ for data, target in test_loader:
+ data, target = data.to(device), target.to(device)
+ output = model(data)
+ test_loss += torch.nn.functional.nll_loss(
+ output, target, reduction="sum"
+ ).item() # sum up batch loss
+ pred = output.argmax(
+ dim=1, keepdim=True
+ ) # get the index of the max log-probability
+ correct += pred.eq(target.view_as(pred)).sum().item()
+
+ test_loss /= len(test_loader.dataset)
+
+ accuracy = 100.0 * correct / len(test_loader.dataset)
+
+ return test_loss, accuracy
+
+ def decode(x):
+ x, _ = torch.max(x, 0)
+ log_p_y = torch.nn.functional.log_softmax(x, dim=1)
+ return log_p_y
+
+ T = 35
+ LR = 0.001
+ EPOCHS = 100 # Increase this for improved accuracy
+
+ if torch.cuda.is_available():
+ DEVICE = torch.device("cuda")
+ else:
+ DEVICE = torch.device("cpu")
+
+ model = Model(
+ encoder=encode.SpikeLatencyLIFEncoder(T), snn=ConvNet(alpha=80), decoder=decode).to(DEVICE)
+
+ optimizer = torch.optim.Adam(model.parameters(), lr=LR)
+
+ # compression
+ if (apply_compression):
+ progressive_compression = ProgressiveCompression(NorseModel=model, maxThreshold=maxTh, alphaP=Alpha, alphaN=-Alpha, to_file=True, apply_reinforcement=reinforcement, file= file)
+
+ training_losses = []
+ mean_losses = []
+ test_losses = []
+ accuracies = []
+
+ for epoch in range(EPOCHS):
+ print(f"Epoch {epoch}")
+ training_loss, mean_loss = train(
+ model, DEVICE, train_loader, optimizer, epoch, max_epochs=EPOCHS
+ )
+ test_loss, accuracy = test(model, DEVICE, test_loader, epoch)
+ training_losses += training_loss
+ mean_losses.append(mean_loss)
+ test_losses.append(test_loss)
+ accuracies.append(accuracy)
+ if (apply_compression):
+ progressive_compression.apply()
+
+ print(f"final accuracy: {accuracies[-1]}")
+ file.write("final accuracy:"+str(accuracies[-1])+"\n")
+ file.write("time:"+str(datetime.now() - before)+"\n")
+
+
+ with open("FashionMNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/FashionMNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before)+".pkl",'wb') as f:
+ torch.save(model,"FashionMNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/FashionMNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before)+".norse")
+ if (apply_compression):
+ pickle.dump([mean_losses,test_losses,accuracies,progressive_compression.weights,progressive_compression.compressions,progressive_compression.thresholds_p,progressive_compression.thresholds_n], f)
+ else:
+ pickle.dump([mean_losses,test_losses,accuracies], f)
diff --git a/Norse/MNIST/MNIST_conv_bench.py b/Norse/MNIST/MNIST_conv_bench.py
new file mode 100644
index 0000000000000000000000000000000000000000..b3188865d6fad069eb0812f3ddeb5e33fc008e66
--- /dev/null
+++ b/Norse/MNIST/MNIST_conv_bench.py
@@ -0,0 +1,215 @@
+import sys
+sys.path.append('../../')
+
+import torch
+import numpy as np
+from compression import ProgressiveCompression
+from norse.torch import LIFCell
+from norse.torch.module.leaky_integrator import LILinearCell
+from norse.torch import LIFParameters
+from norse.torch.module import encode
+#from tqdm import tqdm, trange
+from datetime import datetime
+import torchvision
+import os
+import pickle
+import random
+import itertools
+
+# Reproducibility
+torch.manual_seed(0)
+random.seed(0)
+np.random.seed(0)
+
+MAXTH = [0.6,0.7] #
+ALPHA = [0.005] # [0.002,0.004,0.006,0.008,0.01] #
+REINFORCEMENT = [True] # [False, True]
+
+for maxTh, Alpha, reinforcement in np.array(list(itertools.product(MAXTH, ALPHA, REINFORCEMENT))):
+ try:
+ os.mkdir("MNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement))
+ except OSError as error:
+ print(error)
+
+ for i in range(2):
+ before = datetime.now()
+
+ file = open("MNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/MNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before), 'w+')
+
+ apply_compression = True
+
+ transform = torchvision.transforms.Compose(
+ [
+ torchvision.transforms.ToTensor(),
+ torchvision.transforms.Normalize((0.1307,), (0.3081,)),
+ ]
+ )
+
+ train_data = torchvision.datasets.MNIST(
+ root=".",
+ train=True,
+ download=True,
+ transform=transform,
+ )
+
+ # reduce this number if you run out of GPU memory
+ BATCH_SIZE = 512
+
+ train_loader = torch.utils.data.DataLoader(
+ train_data, batch_size=BATCH_SIZE, shuffle=True
+ )
+
+ test_loader = torch.utils.data.DataLoader(
+ torchvision.datasets.MNIST(
+ root=".",
+ train=False,
+ transform=transform,
+ ),
+ batch_size=BATCH_SIZE,
+ )
+
+ class Model(torch.nn.Module):
+ def __init__(self, encoder, snn, decoder):
+ super(Model, self).__init__()
+ self.encoder = encoder
+ self.snn = snn
+ self.decoder = decoder
+
+ def forward(self, x):
+ x = self.encoder(x)
+ x = self.snn(x)
+ log_p_y = self.decoder(x)
+ return log_p_y
+
+ class ConvNet(torch.nn.Module):
+ def __init__(self, num_channels=1, feature_size=28, method="super", alpha=100):
+ super(ConvNet, self).__init__()
+
+ self.features = int(((feature_size - 4) / 2 - 4) / 2)
+ self.conv1_out_channels = 32
+ self.conv2_out_channels = 128
+ self.fc1_out_channels = 1024
+ self.out_channels = 10
+ self.conv1 = torch.nn.Conv2d(num_channels, self.conv1_out_channels, 5, 1, bias=False)
+ self.conv2 = torch.nn.Conv2d(self.conv1_out_channels, self.conv2_out_channels, 5, 1, bias=False)
+ self.fc1 = torch.nn.Linear(self.features * self.features * self.conv2_out_channels, self.fc1_out_channels, bias=False)
+ self.lif0 = LIFCell(p=LIFParameters(method=method, alpha=alpha,v_th=0.25))
+ self.lif1 = LIFCell(p=LIFParameters(method=method, alpha=alpha,v_th=0.25))
+ self.lif2 = LIFCell(p=LIFParameters(method=method, alpha=alpha,v_th=0.25))
+ self.out = LILinearCell(self.fc1_out_channels, self.out_channels)
+
+ def forward(self, x):
+ seq_length = x.shape[0]
+ batch_size = x.shape[1]
+
+ # specify the initial states
+ s0 = s1 = s2 = so = None
+
+ voltages = torch.zeros(
+ seq_length, batch_size, self.out_channels, device=x.device, dtype=x.dtype
+ )
+
+ for ts in range(seq_length):
+ z = self.conv1(x[ts, :])
+ z, s0 = self.lif0(z, s0)
+ z = torch.nn.functional.max_pool2d(z, 2, 2)
+ z = self.out_channels * self.conv2(z)
+ z, s1 = self.lif1(z, s1)
+ z = torch.nn.functional.max_pool2d(z, 2, 2)
+ z = z.view(-1, 4**2 * self.conv2_out_channels)
+ z = self.fc1(z)
+ z, s2 = self.lif2(z, s2)
+ v, so = self.out(torch.nn.functional.relu(z), so)
+ voltages[ts, :, :] = v
+ return voltages
+
+ def train(model, device, train_loader, optimizer, epoch, max_epochs):
+ model.train()
+ losses = []
+
+ for (data, target) in train_loader: #tqdm(train_loader, leave=False):
+ data, target = data.to(device), target.to(device)
+ optimizer.zero_grad()
+ output = model(data)
+ loss = torch.nn.functional.nll_loss(output, target)
+ loss.backward()
+ optimizer.step()
+ losses.append(loss.item())
+
+ mean_loss = np.mean(losses)
+ return losses, mean_loss
+
+ def test(model, device, test_loader, epoch):
+ model.eval()
+ test_loss = 0
+ correct = 0
+ with torch.no_grad():
+ for data, target in test_loader:
+ data, target = data.to(device), target.to(device)
+ output = model(data)
+ test_loss += torch.nn.functional.nll_loss(
+ output, target, reduction="sum"
+ ).item() # sum up batch loss
+ pred = output.argmax(
+ dim=1, keepdim=True
+ ) # get the index of the max log-probability
+ correct += pred.eq(target.view_as(pred)).sum().item()
+
+ test_loss /= len(test_loader.dataset)
+
+ accuracy = 100.0 * correct / len(test_loader.dataset)
+
+ return test_loss, accuracy
+
+ def decode(x):
+ x, _ = torch.max(x, 0)
+ log_p_y = torch.nn.functional.log_softmax(x, dim=1)
+ return log_p_y
+
+ T = 35
+ LR = 0.001
+ EPOCHS = 100 # Increase this for improved accuracy
+
+ if torch.cuda.is_available():
+ DEVICE = torch.device("cuda:1")
+ else:
+ DEVICE = torch.device("cpu")
+
+ model = Model(
+ encoder=encode.SpikeLatencyLIFEncoder(T), snn=ConvNet(alpha=80), decoder=decode).to(DEVICE)
+
+ optimizer = torch.optim.Adam(model.parameters(), lr=LR)
+
+ # compression
+ if (apply_compression):
+ progressive_compression = ProgressiveCompression(NorseModel=model, maxThreshold=maxTh, alphaP=Alpha, alphaN=-Alpha, to_file=True, apply_reinforcement=reinforcement, file= file)
+
+ training_losses = []
+ mean_losses = []
+ test_losses = []
+ accuracies = []
+
+ for epoch in range(EPOCHS):
+ print(f"Epoch {epoch}")
+ training_loss, mean_loss = train(
+ model, DEVICE, train_loader, optimizer, epoch, max_epochs=EPOCHS
+ )
+ test_loss, accuracy = test(model, DEVICE, test_loader, epoch)
+ training_losses += training_loss
+ mean_losses.append(mean_loss)
+ test_losses.append(test_loss)
+ accuracies.append(accuracy)
+ if (apply_compression):
+ progressive_compression.apply()
+
+ print(f"final accuracy: {accuracies[-1]}")
+ file.write("final accuracy:"+str(accuracies[-1])+"\n")
+ file.write("time:"+str(datetime.now() - before)+"\n")
+
+
+ with open("MNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/MNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before)+".pkl",'wb') as f:
+ torch.save(model,"MNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"/MNIST_CONV_maxTh:"+str(maxTh)+"_Alpha:"+str(Alpha)+"_"+"reinforcement:"+str(reinforcement)+"_"+str(before)+".norse")
+ if (apply_compression):
+ pickle.dump([mean_losses,test_losses,accuracies,progressive_compression.weights,progressive_compression.compressions,progressive_compression.thresholds_p,progressive_compression.thresholds_n], f)
+ else:
+ pickle.dump([mean_losses,test_losses,accuracies], f)
diff --git a/Norse/compression.py b/Norse/compression.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1e169d3404a6e7583ea077a72bb8fc9a1c62f03
--- /dev/null
+++ b/Norse/compression.py
@@ -0,0 +1,101 @@
+import numpy as np
+import torch
+
+class ProgressiveCompression:
+ def __init__(self, NorseModel,maxThreshold=0.3, alphaP=0.005, alphaN=-0.005, betaP=0.01, betaN=-0.01, prune_recurrent=True, apply_reinforcement=False, to_file=False, file=None, layerwise=False):
+ self.alpha_p = []
+ self.alpha_n = []
+ self.max_threshold = maxThreshold
+ self.model = NorseModel
+ self.prune_recurrent = prune_recurrent
+ self.layerwise = layerwise
+ self.apply_reinforcement = apply_reinforcement
+ self.to_file = to_file
+ self.file = file
+ self.max_threshold_p = []
+ self.max_threshold_n = []
+ self.max_w_p = []
+ self.max_w_n = []
+ self.prune_matrix = []
+ self.thresholds_p = []
+ self.thresholds_n = []
+ self.compressions = []
+ self.weights = []
+ self.betaP = betaP
+ self.betaN = betaN
+
+ i = 0
+ for name, param in self.model.named_parameters():
+ print(name,param.data.min(),param.data.max())
+ self.max_w_p.append(param.data.max())
+ self.max_w_n.append(param.data.min())
+
+ if (self.layerwise):
+ self.alpha_p.append(alphaP + (0.005*i))
+ self.alpha_n.append(alphaN + (-0.005*i))
+ else:
+ self.alpha_p.append(alphaP)
+ self.alpha_n.append(alphaN)
+
+ self.max_threshold_p.append(param.data.max() * self.max_threshold)
+ self.max_threshold_n.append(param.data.min() * self.max_threshold)
+ print("max_threshold_n: "+str(param.data.min() * self.max_threshold)+" max_threshold_p: "+str(param.data.max() * self.max_threshold)+"\n")
+ if(self.to_file):
+ self.file.write(str(name)+" "+str(param.data.min())+" "+str(param.data.max())+"\n")
+ self.file.write("max_threshold_n: "+str(param.data.min() * self.max_threshold)+" max_threshold_p: "+str(param.data.max() * self.max_threshold)+"\n")
+ i = i + 1
+
+ for param in self.model.parameters():
+ self.prune_matrix.append(torch.zeros_like(param))
+ self.thresholds_p.append([None])
+ self.thresholds_n.append([None])
+ self.compressions.append([0])
+
+ def applyprune(self,name, alpha_p, alpha_n, max_thresholds_p,max_thresholds_n, weights, prune_matrix, threshold_p, threshold_n):
+ print(name, " before prune: Min: ",weights.min()," Max: ",weights.max())
+ if ((not self.prune_recurrent) and ("recurrent" in name)):
+ return weights, prune_matrix, threshold_p, threshold_n
+ else:
+ if threshold_p == None:
+ threshold_p = alpha_p
+ else:
+ if threshold_p < max_thresholds_p:
+ threshold_p += alpha_p * (np.count_nonzero(prune_matrix.cpu() == 0) / (np.count_nonzero(prune_matrix.cpu() == 0) + np.count_nonzero(prune_matrix.cpu() == 1)))
+ threshold_p = round(threshold_p, 6)
+
+ if threshold_n == None:
+ threshold_n = alpha_n
+ else:
+ if threshold_n > max_thresholds_n:
+ threshold_n += alpha_n * (np.count_nonzero(prune_matrix.cpu() == 0) / (np.count_nonzero(prune_matrix.cpu() == 0) + np.count_nonzero(prune_matrix.cpu() == 1)))
+ threshold_n = round(threshold_n, 6)
+
+ return weights.masked_fill(((weights < threshold_p) & (weights > 0)) | ((weights > threshold_n) & (weights < 0)), 0), prune_matrix.masked_fill(((weights < threshold_p) & (weights > 0)) | ((weights > threshold_n) & (weights < 0)), 1), threshold_p, threshold_n
+
+ def reinforcement(self, name, weights, beta_p, beta_n, thres_p, thres_n, max_w_p, max_w_n):
+ if ((not self.prune_recurrent) and ("recurrent" in name)):
+ return weights
+ weights = torch.where(weights > 0,weights + beta_p * thres_p,weights)
+ weights[weights > max_w_p] = max_w_p
+ weights = torch.where(weights < 0,weights - beta_n * thres_n,weights)
+ weights[weights < max_w_n] = max_w_n
+ return weights
+
+ def apply(self,):
+ i = 0
+ for name, param in self.model.named_parameters():
+ param.data, self.prune_matrix[i], thres_p , thres_n = self.applyprune(name, self.alpha_p[i],self.alpha_n[i],self.max_threshold_p[i],self.max_threshold_n[i],param.data,self.prune_matrix[i],self.thresholds_p[i][-1],self.thresholds_n[i][-1])
+ #try:
+ print(name,"zeros:",int((self.prune_matrix[i]).sum()),"/",float(torch.prod(torch.tensor(param.data.shape))),"("+str(round(float((self.prune_matrix[i]).sum()*100/(torch.prod(torch.tensor(param.data.shape)))),3))+"%)","threshold_n:",thres_n,"threshold_p:",thres_p)
+ if(self.to_file):
+ self.file.write(str(name)+" zeros: "+str(float((self.prune_matrix[i]).sum()))+" / "+str(float(torch.prod(torch.tensor(param.data.shape))))+" ("+str(round(float((self.prune_matrix[i]).sum()*100/(torch.prod(torch.tensor(param.data.shape)))),3))+"%) threshold_n: "+str(thres_n)+" threshold_p: "+str(thres_p)+"\n")
+ self.compressions[i].append(round(float((self.prune_matrix[i]).sum()*100/(torch.prod(torch.tensor(param.data.shape)))),3))
+ self.thresholds_p[i].append(thres_p)
+ self.thresholds_n[i].append(thres_n)
+ if(self.apply_reinforcement):
+ param.data = self.reinforcement(name, param.data, self.betaP, self.betaN, thres_p, thres_n, self.max_w_p[i], self.max_w_n[i])
+ self.weights.append(param.data)
+ #except Exception:
+ # pass
+
+ i+=1