diff options
Diffstat (limited to 'neurobench_testing/memTorch_testing.py')
| -rw-r--r-- | neurobench_testing/memTorch_testing.py | 81 |
1 files changed, 67 insertions, 14 deletions
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 |