From 30dc9553e8a7a7d603077a80238782298d0b9349 Mon Sep 17 00:00:00 2001
From: Tanner Robison <Robison.tnr@gmail.com>
Date: Fri, 19 Jun 2026 00:02:30 -0700
Subject: customizeable memristor model

started trying to implement a customizeable memristor model to pass to memTorch, porting a model I got from kaevin too fit the pyTorch framework
---
 neurobench_testing/memTorch_testing.py | 81 ++++++++++++++++++++++++++++------
 1 file changed, 67 insertions(+), 14 deletions(-)

(limited to 'neurobench_testing/memTorch_testing.py')

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
-- 
cgit v1.2.3