Skip to content

Saving and Loading Models

nwave models are standard PyTorch nn.Module objects, so saving and loading follow the usual PyTorch patterns.

One rule to remember: always call prepare_net(model) after loading weights. It resets the internal states (membrane potentials, spike history) to zero so the network is in a clean state, ready to process a new sequence.

import torch
import torch.nn as nn
from nwavesdk.layers import H1v2Synapse, H1v2Layer, prepare_net

N_IN, N_OUT, DT = 16, 8, 1e-3

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.syn = H1v2Synapse(N_IN, N_OUT, device="cpu")
        self.lyr = H1v2Layer(N_OUT, taus=10e-3, dt=DT, layer_topology="FF", device="cpu")

    def forward(self, x):
        prepare_net(self)
        spk_list = []
        for t in range(x.shape[1]):
            cur = self.syn(x[:, t, :])
            spk, _ = self.lyr(cur)
            spk_list.append(spk)
        return torch.stack(spk_list, dim=1)

model = SimpleNet()
model.eval()
print("Model created.")

Model created.

# --- Save ---
# The recommended approach: save the state dict (weights only, not the class itself).
# This is portable across code refactors and Python versions.
torch.save(model.state_dict(), "/tmp/my_nwave_model.pt")
print("Saved to /tmp/my_nwave_model.pt")

Saved to /tmp/my_nwave_model.pt

# --- Load ---
# 1. Create a fresh instance of the same architecture.
loaded_model = SimpleNet()

# 2. Restore the weights.
state = torch.load("/tmp/my_nwave_model.pt", map_location="cpu")
loaded_model.load_state_dict(state)

# 3. Always call prepare_net after loading — it resets neuron states to zero.
loaded_model.eval()
print("Model loaded.")

Model loaded.

# Verify: both models produce identical output on the same input.
x = (torch.rand(2, 50, N_IN) > 0.8).float()

with torch.no_grad():
    out_original = model(x)
    out_loaded   = loaded_model(x)

match = torch.allclose(out_original, out_loaded)
print(f"Outputs match: {match}")

Outputs match: True

Saving training state

To resume training from a checkpoint, save both the model weights and the optimizer state:

torch.save({
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "epoch": current_epoch,
}, "checkpoint.pt")

# Resume
checkpoint = torch.load("checkpoint.pt", map_location="cpu")
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
start_epoch = checkpoint["epoch"] + 1

Further reading