NIR N-MNIST SCNN

This example imports a convolutional SNN from NIR (scnn_mnist.nir, trained in snnTorch) with Spyx's run-and-return from_nir and reproduces the first conv+IF layer's activity. The final accuracy section downloads the ~1 GB N-MNIST test set on first run.

In [1]:

import jax.numpy as jnp
import matplotlib.pyplot as plt
import nir
import numpy as np

import spyx
import spyx.nn

import jax.numpy as jnp import matplotlib.pyplot as plt import nir import numpy as np import spyx import spyx.nn

In [ ]:

# val_numbers is (B, T, C, H, W). Spyx conv is channels-last, so make it
# (T, B, H, W, C): time-major for snn.run, channels last for nnx.Conv.
input_data = jnp.transpose(
    jnp.array(np.load("val_numbers.npy"), dtype=jnp.float32), axes=(1, 0, 3, 4, 2)
)

# val_numbers is (B, T, C, H, W). Spyx conv is channels-last, so make it # (T, B, H, W, C): time-major for snn.run, channels last for nnx.Conv. input_data = jnp.transpose( jnp.array(np.load("val_numbers.npy"), dtype=jnp.float32), axes=(1, 0, 3, 4, 2) )

In [3]:

input_data.shape # had to transpose data because it's stored in time-major format...

input_data.shape # had to transpose data because it's stored in time-major format...

Out[3]:

(10, 300, 34, 34, 2)

In [ ]:

ng = nir.read("scnn_mnist.nir")

ng = nir.read("scnn_mnist.nir")

In [ ]:

# Grab the first conv + spiking layer as a standalone subgraph. The Output
# node must declare node "1"'s (first IF) shape, not the 10-class network output.
first_conv = nir.NIRGraph(
    nodes={
        "input": ng.nodes["input"],
        "0": ng.nodes["0"],
        "1": ng.nodes["1"],
        "output": nir.Output(output_type={"output": np.array([16, 16, 16])}),
    },
    edges=[("input", "0"), ("0", "1"), ("1", "output")],
)

# Grab the first conv + spiking layer as a standalone subgraph. The Output # node must declare node "1"'s (first IF) shape, not the 10-class network output. first_conv = nir.NIRGraph( nodes={ "input": ng.nodes["input"], "0": ng.nodes["0"], "1": ng.nodes["1"], "output": nir.Output(output_type={"output": np.array([16, 16, 16])}), }, edges=[("input", "0"), ("0", "1"), ("1", "output")], )

In [ ]:

# from_nir now builds AND runs the model, returning (model, (outputs, states)).
# return_all_states gives the per-timestep neuron states (membrane traces).
fl_SNN, (output_spikes, membrane_potentials) = spyx.nir.from_nir(
    first_conv, input_data, dt=1, return_all_states=True
)

# from_nir now builds AND runs the model, returning (model, (outputs, states)). # return_all_states gives the per-timestep neuron states (membrane traces). fl_SNN, (output_spikes, membrane_potentials) = spyx.nir.from_nir( first_conv, input_data, dt=1, return_all_states=True )

In [7]:

output_spikes.shape  # (T, B, H, W, C) spike train from the run

output_spikes.shape # (T, B, H, W, C) spike train from the run

Out[7]:

(10, 300, 16, 16, 16)

In [8]:

save_data = np.array(jnp.transpose(output_spikes, (1,0,2,3,4)))
save_data.shape

save_data = np.array(jnp.transpose(output_spikes, (1,0,2,3,4))) save_data.shape

Out[8]:

(300, 10, 16, 16, 16)

In [ ]:

spyx_act = jnp.sum(output_spikes[0], axis=0)

spyx_act = jnp.sum(output_spikes[0], axis=0)

In [10]:

spyx_act.shape

spyx_act.shape

Out[10]:

(16, 16, 16)

In [11]:

plt.imshow(jnp.sum(spyx_act, axis=(0)))
plt.colorbar()
plt.title("Spyx SCNN Spiking, first Conv+IF Layer")
plt.show()

plt.imshow(jnp.sum(spyx_act, axis=(0))) plt.colorbar() plt.title("Spyx SCNN Spiking, first Conv+IF Layer") plt.show()

In [ ]:

import numpy as np

snntorch_act = np.sum(np.load("./snnTorch_activity.npy", "r"), axis=0)

import numpy as np snntorch_act = np.sum(np.load("./snnTorch_activity.npy", "r"), axis=0)

In [13]:

plt.imshow(jnp.sum(snntorch_act[0], axis=(0)))
plt.colorbar()
plt.title("snnTorch SCNN First Conv+IF Layer")
plt.show()

plt.imshow(jnp.sum(snntorch_act[0], axis=(0))) plt.colorbar() plt.title("snnTorch SCNN First Conv+IF Layer") plt.show()

In [ ]:

def cosine_similarity(vector1, vector2):
    # Compute the dot product of the two vectors
    dot_product = jnp.dot(vector1, vector2)

    # Compute the magnitude (L2 norm) of each vector
    magnitude1 = jnp.linalg.norm(vector1)
    magnitude2 = jnp.linalg.norm(vector2)

    # Compute the cosine similarity
    similarity = dot_product / (magnitude1 * magnitude2)

    return similarity

def cosine_similarity(vector1, vector2): # Compute the dot product of the two vectors dot_product = jnp.dot(vector1, vector2) # Compute the magnitude (L2 norm) of each vector magnitude1 = jnp.linalg.norm(vector1) magnitude2 = jnp.linalg.norm(vector2) # Compute the cosine similarity similarity = dot_product / (magnitude1 * magnitude2) return similarity

In [15]:

cosine_similarity(jnp.sum(spyx_act, 0).flatten(), jnp.sum(snntorch_act[0], 0).flatten())

cosine_similarity(jnp.sum(spyx_act, 0).flatten(), jnp.sum(snntorch_act[0], 0).flatten())

Out[15]:

Array(nan, dtype=float32)

In [ ]:

np.save("spyx_activity.npy", save_data)

np.save("spyx_activity.npy", save_data)

Inference Accuracy¶

In [ ]:

# NOTE: downloads the ~1 GB N-MNIST test set on first run.
import tonic
import torch

bs = 128
collate = tonic.collation.PadTensors(batch_first=False)
to_frame = tonic.transforms.ToFrame(sensor_size=tonic.datasets.NMNIST.sensor_size, time_window=1e3)
test_ds = tonic.datasets.NMNIST("./nmnist", transform=to_frame, train=False)
test_dl = torch.utils.data.DataLoader(test_ds, shuffle=True, batch_size=bs, collate_fn=collate)

# NOTE: downloads the ~1 GB N-MNIST test set on first run. import tonic import torch bs = 128 collate = tonic.collation.PadTensors(batch_first=False) to_frame = tonic.transforms.ToFrame(sensor_size=tonic.datasets.NMNIST.sensor_size, time_window=1e3) test_ds = tonic.datasets.NMNIST("./nmnist", transform=to_frame, train=False) test_dl = torch.utils.data.DataLoader(test_ds, shuffle=True, batch_size=bs, collate_fn=collate)

In [17]:

# Import + run the full SCNN. from_nir reconstructs the whole graph (convs,
# spiking layers, SumPool, flatten, dense) and runs it on input_data.
SNN, full_out = spyx.nir.from_nir(ng, input_data, dt=1)
print("full SCNN output:", full_out.shape)  # (T, B, 10)

# Import + run the full SCNN. from_nir reconstructs the whole graph (convs, # spiking layers, SumPool, flatten, dense) and runs it on input_data. SNN, full_out = spyx.nir.from_nir(ng, input_data, dt=1) print("full SCNN output:", full_out.shape) # (T, B, 10)

full SCNN output: (10, 300, 10)

In [ ]:

# Tonic frames are (T, B, C, H, W); make them channels-last for the spyx conv.
accs = []
for x, y in test_dl:
    x = jnp.transpose(jnp.array(x, dtype=jnp.float32), (0, 1, 3, 4, 2))
    spikes, _ = spyx.nn.run(SNN, x)
    acc, preds = spyx.fn.integral_accuracy(
        jnp.transpose(spikes, (1, 0, 2)), jnp.array(y)
    )
    accs.append(acc)

# Tonic frames are (T, B, C, H, W); make them channels-last for the spyx conv. accs = [] for x, y in test_dl: x = jnp.transpose(jnp.array(x, dtype=jnp.float32), (0, 1, 3, 4, 2)) spikes, _ = spyx.nn.run(SNN, x) acc, preds = spyx.fn.integral_accuracy( jnp.transpose(spikes, (1, 0, 2)), jnp.array(y) ) accs.append(acc)

In [ ]:

final_acc = np.mean(np.array(accs))
final_acc

final_acc = np.mean(np.array(accs)) final_acc

In [ ]:

np.save("spyx_accuracy.npy", final_acc)

np.save("spyx_accuracy.npy", final_acc)