diff options
| -rw-r--r-- | neurobench_testing/custom_memristor_model.py | 42 | ||||
| -rw-r--r-- | neurobench_testing/memTorch_testing.py | 81 | ||||
| -rw-r--r-- | neurobench_testing/memTorch_testing_custom_memristor.py | 160 | ||||
| -rw-r--r-- | references/kaevin_memristor.py | 134 |
4 files changed, 403 insertions, 14 deletions
diff --git a/neurobench_testing/custom_memristor_model.py b/neurobench_testing/custom_memristor_model.py new file mode 100644 index 0000000..8bdcde0 --- /dev/null +++ b/neurobench_testing/custom_memristor_model.py @@ -0,0 +1,42 @@ +import torch +from memtorch.bh.memristor.Memristor import Memristor + +class MemtorchMemristor(Memristor): + def __init__( + self, + k_off = 1.0, # switching rate for off state + k_on = -1.0, # switching rate for on state + alpha_off = 5, # exponent controlling nonlinearity + alpha_on = 5, # exponent controlling nonlinearity + i_off = 0.5e-3, # threshhold current to trigger off state + i_on = 0.5e-3, # threshold current to trigger on state + r_on = 1e3, # maximum resistance + r_off = 10e3, # minimum resistance + p = 2, # window function exponent + **kwargs + ): + #initializing base memristor class + super(MemtorchMemristor, self).__init__(r_off=r_off, r_on=r_on, **kwargs) + + # hyper parameters + self.k_off = k_off + self.k_on = k_on + self.alpha_off = alpha_off + self.alpha_on = alpha_on + self.i_on = i_on + self.i_off = i_off + self.p = p + + # makes sure w starts in valid state + if not hasattr(self, 'w'): + self.w = torch.tensor(0.5) + + + """ + Updates w and computes new resistance + + """ + def step(self, v, dt): + i = v / self.r_curr + + diff --git a/neurobench_testing/memTorch_testing.py b/neurobench_testing/memTorch_testing.py index a44fd98..5c3d828 100644 --- a/neurobench_testing/memTorch_testing.py +++ b/neurobench_testing/memTorch_testing.py @@ -1,14 +1,19 @@ -from memtorch import memristor +""" +This is a program to test out using memTorch with neurobench. +June 17th, 2026 +Author: Tanner Robison, +Teuscher Lab +""" import torch import torch.nn as nn import snntorch as snn +from snntorch import surrogate from torch.utils.data import DataLoader from neurobench.models import SNNTorchModel from neurobench.benchmarks import Benchmark, benchmark from neurobench.datasets import SpeechCommands - from neurobench.metrics.workload import ( ActivationSparsity, SynapticOperations, @@ -29,8 +34,7 @@ from memtorch.map.Parameter import naive_map from memtorch.bh.memristor import VTEAM beta = 0.9 - -class SimpleSNN(nn.Module): +class SNN(nn.Module): def __init__(self): super().__init__() @@ -54,36 +58,46 @@ class SimpleSNN(nn.Module): return spk2, mem2 device = torch.device("cpu") -net = SimpleSNN().to(device) +spike_grad = surrogate.fast_sigmoid() +net = nn.Sequential( + nn.Flatten(), + nn.Linear(20, 256), + snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True), + nn.Linear(256, 256), + snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True), + nn.Linear(256, 256), + snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True), + nn.Linear(256, 35), + snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True), +) #memristor patch reference_memristor = VTEAM() -print("Patching model to memristive crossbar") +net.load_state_dict(torch.load("examples/gsc/model_data/s2s_gsc_snntorch", map_location=device)) + patched_net = patch_model( copy.deepcopy(net), memristor_model=reference_memristor, - memristor_model_params={}, + memristor_model_params={'time_series_resolution': 1e-8}, mapping_routine=naive_map, transistor=True, + tile_shape=(128, 128), ADC_resolution=8, - use_bindings=False + use_bindings=True ) -model = SNNTorchModel(patched_net) - static_metrics = [Footprint, ConnectionSparsity] workload_metrics = [ActivationSparsity, SynapticOperations, ClassificationAccuracy] # data loader here maybe?? test_set = SpeechCommands(path="data/SpeechCommands/", subset="testing") -test_set_loader = DataLoader(test_set, batch_size=50, shuffle=True) +test_set_loader = DataLoader(test_set, batch_size=500, shuffle=True) pre_processor = [S2SPreProcessor(device=device)] post_processor = [ChooseMaxCount()] -# print("Layer 1 Beta:", net[2].beta) - +model = SNNTorchModel(net) benchmark = Benchmark( model, test_set_loader, @@ -93,7 +107,46 @@ benchmark = Benchmark( ) results = benchmark.run() -print(results) +print("\n\n----- IDEAL BENCHMARK -----") +print(f"Footprint: {results['Footprint']}") +print(f"Connection Sparsity: {results['ConnectionSparsity']}") +print(f"Activation Sparsity: {results['ActivationSparsity']}") +print(f"Synaptic Operations: {results['SynapticOperations']}") +print(f"Classification Accuracy: {results['ClassificationAccuracy']}") + +#energy calculations +ENERGY_PER_MAC = 0.9e-12 +ENERGY_PER_AC = 0.1e-12 + +macs = results['SynapticOperations']['Effective_MACs'] +acs = results['SynapticOperations']['Effective_ACs'] + +total_energy = (macs * ENERGY_PER_MAC) + (acs * ENERGY_PER_AC) + +print("Energy Report:") +print(f"Total Operations: {macs} MACS, {acs} acs ") +print(f"Calculated Energy cost: {total_energy} joules per batch") + + +model = SNNTorchModel(patched_net) +benchmark = Benchmark( + model, + test_set_loader, + pre_processor, + post_processor, + [static_metrics, workload_metrics] +) + +with torch.no_grad(): #makes sure its in inference mode + #otherwise you get memory leaks + results = benchmark.run() + +print("\n\n----- MEMRISTOR BENCHMARK -----") +print(f"Footprint: {results['Footprint']}") +print(f"Connection Sparsity: {results['ConnectionSparsity']}") +print(f"Activation Sparsity: {results['ActivationSparsity']}") +print(f"Synaptic Operations: {results['SynapticOperations']}") +print(f"Classification Accuracy: {results['ClassificationAccuracy']}") #energy calculations ENERGY_PER_MAC = 0.9e-12 diff --git a/neurobench_testing/memTorch_testing_custom_memristor.py b/neurobench_testing/memTorch_testing_custom_memristor.py new file mode 100644 index 0000000..64f338e --- /dev/null +++ b/neurobench_testing/memTorch_testing_custom_memristor.py @@ -0,0 +1,160 @@ +from memtorch import memristor +import torch +import torch.nn as nn +import snntorch as snn +from snntorch import surrogate + +from torch.utils.data import DataLoader + +from neurobench.models import SNNTorchModel +from neurobench.benchmarks import Benchmark, benchmark +from neurobench.datasets import SpeechCommands + +from neurobench.metrics.workload import ( + ActivationSparsity, + SynapticOperations, + ClassificationAccuracy, +) + +from neurobench.metrics.static import ( + Footprint, + ConnectionSparsity, +) + +from neurobench.processors.preprocessors import S2SPreProcessor +from neurobench.processors.postprocessors import ChooseMaxCount + +import copy +from memtorch.mn.Module import patch_model +from memtorch.map.Parameter import naive_map +from memtorch.bh.memristor import VTEAM + +from kaevin_memristor import TEAMMemristor + +beta = 0.9 +class SNN(nn.Module): + def __init__(self): + super().__init__() + + #standard layers + self.fc1 = nn.Linear(20, 128) + self.fc2 = nn.Linear(128, 35) + + #Spiking neurons + self.lif1 = snn.Leaky(beta=beta, init_hidden=True) + self.lif2 = snn.Leaky(beta=beta, init_hidden=True, output=True) + + def forward(self, x): + x = x.view(x.size(0), -1) + + cur1 = self.fc1(x) + spk1 = self.lif1(cur1) + + cur2 = self.fc2(spk1) + spk2, mem2 = self.lif2(cur2) + + return spk2, mem2 + +device = torch.device("cpu") +spike_grad = surrogate.fast_sigmoid() +net = nn.Sequential( + nn.Flatten(), + nn.Linear(20, 256), + snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True), + nn.Linear(256, 256), + snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True), + nn.Linear(256, 256), + snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True), + nn.Linear(256, 35), + snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True), +) + +#memristor patch +reference_memristor = TEAMMemristor + +net.load_state_dict(torch.load("examples/gsc/model_data/s2s_gsc_snntorch", map_location=device)) + +patched_net = patch_model( + copy.deepcopy(net), + memristor_model=reference_memristor, + memristor_model_params={'time_series_resolution': 1e-8}, + mapping_routine=naive_map, + transistor=True, + tile_shape=(128, 128), + ADC_resolution=8, + use_bindings=True +) + +static_metrics = [Footprint, ConnectionSparsity] +workload_metrics = [ActivationSparsity, SynapticOperations, ClassificationAccuracy] + +# data loader here maybe?? +test_set = SpeechCommands(path="data/SpeechCommands/", subset="testing") +test_set_loader = DataLoader(test_set, batch_size=500, shuffle=True) + +pre_processor = [S2SPreProcessor(device=device)] +post_processor = [ChooseMaxCount()] + +model = SNNTorchModel(net) +benchmark = Benchmark( + model, + test_set_loader, + pre_processor, + post_processor, + [static_metrics, workload_metrics] +) + +results = benchmark.run() +print("\n\n----- IDEAL BENCHMARK -----") +print(f"Footprint: {results['Footprint']}") +print(f"Connection Sparsity: {results['ConnectionSparsity']}") +print(f"Activation Sparsity: {results['ActivationSparsity']}") +print(f"Synaptic Operations: {results['SynapticOperations']}") +print(f"Classification Accuracy: {results['ClassificationAccuracy']}") + +#energy calculations +ENERGY_PER_MAC = 0.9e-12 +ENERGY_PER_AC = 0.1e-12 + +macs = results['SynapticOperations']['Effective_MACs'] +acs = results['SynapticOperations']['Effective_ACs'] + +total_energy = (macs * ENERGY_PER_MAC) + (acs * ENERGY_PER_AC) + +print("Energy Report:") +print(f"Total Operations: {macs} MACS, {acs} acs ") +print(f"Calculated Energy cost: {total_energy} joules per batch") + + +model = SNNTorchModel(patched_net) +benchmark = Benchmark( + model, + test_set_loader, + pre_processor, + post_processor, + [static_metrics, workload_metrics] +) + +with torch.no_grad(): #makes sure its in inference mode + #otherwise you get memory leaks + results = benchmark.run() + +print("\n\n----- MEMRISTOR BENCHMARK -----") +print(f"Footprint: {results['Footprint']}") +print(f"Connection Sparsity: {results['ConnectionSparsity']}") +print(f"Activation Sparsity: {results['ActivationSparsity']}") +print(f"Synaptic Operations: {results['SynapticOperations']}") +print(f"Classification Accuracy: {results['ClassificationAccuracy']}") + +#energy calculations +ENERGY_PER_MAC = 0.9e-12 +ENERGY_PER_AC = 0.1e-12 + +macs = results['SynapticOperations']['Effective_MACs'] +acs = results['SynapticOperations']['Effective_ACs'] + +total_energy = (macs * ENERGY_PER_MAC) + (acs * ENERGY_PER_AC) + +print("Energy Report:") +print(f"Total Operations: {macs} MACS, {acs} acs ") +print(f"Calculated Energy cost: {total_energy} joules per batch") diff --git a/references/kaevin_memristor.py b/references/kaevin_memristor.py new file mode 100644 index 0000000..b4a4364 --- /dev/null +++ b/references/kaevin_memristor.py @@ -0,0 +1,134 @@ +import numpy as np +from scipy.integrate import solve_ivp + +# TEAM (Threshold Adaptive Memristor) model as a reusable class +class TEAMMemristor: + + def __init__( + self, + k_off=1, # switching rate for off-state + k_on=-1, # switching rate for on-state + alpha_off=5, # exponent controlling nonlinearity when switching off + alpha_on=5, # exponent controlling nonlinearity when switching on + i_off=0.5e-3, # threshold current to trigger off state switching + i_on=-0.5e-3, # threshold current to trigger on state switching + g_on=1/1e3, # maximum conductance (1 kohm = 1 ms) + g_off=1/10e3, # minimum conductance (10 kohm = 0.1 ms) + w_init=0.5, # initial state variable (0=off, 1=on) + p=2 # window function exponent + ): + self.k_off = k_off + self.k_on = k_on + self.alpha_off = alpha_off + self.alpha_on = alpha_on + self.i_off = i_off + self.i_on = i_on + self.G_on = G_on + self.G_off = G_off + self.w_init = w_init + self.p = p + + def set_state(self, w): + # Update initial condition for next simulation + self.w_init = np.clip(w, 0, 1) + + def window(self, w, i): + # Nonlinear window function: reduces switching rate near the boundaries + w = np.clip(w, 0.0, 1.0) + if i >= 0: + return 1 - w**(2*self.p) # switching off: slower near w=1 + else: + return 1 - (1-w)**(2*self.p) # switching on: slower near w=0 + + def conductance(self, w): + # Linear interpolation between off and on conductance based on state w + w = np.clip(w, 0.0, 1.0) + return self.G_off + w*(self.G_on - self.G_off) + + def dw_dt(self, w, i): + # TEAM state dynamics: dw/dt depends on current magnitude and direction + w = np.clip(w, 0.0, 1.0) + + if i >= self.i_off: # positive current above threshold = switch off + dw = ( + self.k_off + * ((i/self.i_off)-1)**self.alpha_off + * self.window(w, i) + ) + elif i <= self.i_on: # negative current below threshold = switch on + dw = ( + self.k_on + * (((-i)/abs(self.i_on))-1)**self.alpha_on + * self.window(w, i) + ) + else: # between thresholds = no switching + dw = 0.0 + + # Enforce physical bounds: prevent state from leaving [0,1] + if w <= 0 and dw < 0: + dw = 0 + if w >= 1 and dw > 0: + dw = 0 + + return dw + + def simulate(self, + freq=1, # excitation frequency (Hz) + V_amp=1.5, # sinusoid amplitude (V) + cycles=3): # number of periods to simulate + # Solve the memristor ODE for given frequency and voltage amplitude + + def voltage(t): # sinusoidal excitation signal + return V_amp*np.sin(2*np.pi*freq*t) + + def ode(t, y): # dy/dt: current through memristor + w = y[0] + v = voltage(t) + G = self.conductance(w) + i = G*v # Ohm's law: i = G*v + return [self.dw_dt(w, i)] + + T = 1/freq # period + t_end = cycles*T + t_eval = np.linspace(0, t_end, 10000) # dense time grid for smooth curves + + # Solve with RK45, small max step + sol = solve_ivp( + ode, + [0, t_end], + [self.w_init], + t_eval=t_eval, + method='RK45', + max_step=T/1000, # max step keeps resolution within one period + rtol=1e-8, + atol=1e-10 + ) + + raw_w = sol.y[0] + + # Check if numerical solver violated physical bounds + eps = 1e-6 + if np.any(raw_w < -eps) or np.any(raw_w > 1+eps): + print( + "WARNING: solver left bounds " + f"min={raw_w.min():.12f}, " + f"max={raw_w.max():.12f}" + ) + + w = np.clip(raw_w, 0, 1) # enforce bounds just in case + t = sol.t + v = voltage(t) + G = self.conductance(w) + i = G*v # current throughout simulation + + return t, w, v, i + + def resistance(self, w): + # Compute resistance as reciprocal of conductance + return 1/self.conductance(w) + + def reset(self): + # Reset state to default initial condition + self.w_init = 0.5 + + |