diff options
Diffstat (limited to 'neurobench_testing/memTorch_testing.py')
| -rw-r--r-- | neurobench_testing/memTorch_testing.py | 89 |
1 files changed, 48 insertions, 41 deletions
diff --git a/neurobench_testing/memTorch_testing.py b/neurobench_testing/memTorch_testing.py index 5c3d828..1992e05 100644 --- a/neurobench_testing/memTorch_testing.py +++ b/neurobench_testing/memTorch_testing.py @@ -1,15 +1,10 @@ -""" -This is a program to test out using memTorch with neurobench. -June 17th, 2026 -Author: Tanner Robison, -Teuscher Lab -""" +import sys import torch import torch.nn as nn import snntorch as snn from snntorch import surrogate -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, Subset from neurobench.models import SNNTorchModel from neurobench.benchmarks import Benchmark, benchmark @@ -32,31 +27,9 @@ import copy from memtorch.mn.Module import patch_model from memtorch.map.Parameter import naive_map from memtorch.bh.memristor import VTEAM +from memtorch.map.Input import naive_scale 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( @@ -71,28 +44,38 @@ net = nn.Sequential( snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True), ) -#memristor patch -reference_memristor = VTEAM() +vteam_params = { + 'time_series_resolution': 1e-3, + 'r_on': 50, + 'r_off': 1000, +} 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}, + memristor_model=VTEAM, + memristor_model_params=vteam_params, + module_parameters_to_patch=[torch.nn.Linear], mapping_routine=naive_map, transistor=True, tile_shape=(128, 128), - ADC_resolution=8, - use_bindings=True + max_input_voltage=0.3, + scaling_routine=naive_scale, + ADC_resolution=16, + use_bindings=True, + verbose=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) + +#shorten data set so I can actually run it lol +tiny_indices = list(range(10)) +tiny_test_set = Subset(test_set, tiny_indices) +test_set_loader = DataLoader(tiny_test_set, batch_size=16, shuffle=True) pre_processor = [S2SPreProcessor(device=device)] post_processor = [ChooseMaxCount()] @@ -106,13 +89,36 @@ benchmark = Benchmark( [static_metrics, workload_metrics] ) +print("\n --Checking signal strength--") +dummy_input = torch.randn(2, 20).to(device) + +try: + raw_signal = patched_net(dummy_input) + + if isinstance(raw_signal, tuple) and len(raw_signal) > 1: + voltage_signal = raw_signal[0] + print("Checking Neuron voltages") + else: + voltage_signal = raw_signal + print("Checking RAW OUTPUTS") + + + print(f"Signal Max: {raw_signal.max().item():.8f}") + print(f"Signal Min: {raw_signal.min().item():.8f}") + print(f"Signal Mean: {raw_signal.mean().item():.8f}") + +except Exception as e: + print("Error getting signal:", e) + +sys.exit() + 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']}") +print(f"Classification Accuracy: {results['ClassificationAccuracy']}\n") #energy calculations ENERGY_PER_MAC = 0.9e-12 @@ -125,7 +131,7 @@ 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") +print(f"Calculated Energy cost: {total_energy} joules per batch\n\n") model = SNNTorchModel(patched_net) @@ -137,8 +143,9 @@ benchmark = Benchmark( [static_metrics, workload_metrics] ) + with torch.no_grad(): #makes sure its in inference mode - #otherwise you get memory leaks + #otherwise you get memory leaks : ( results = benchmark.run() print("\n\n----- MEMRISTOR BENCHMARK -----") |