def __init__(self,
device,
num_channels=1,
feature_size=28,
method="super",
dtype=torch.float):
super(ConvNet4, self).__init__()
self.features = int(((feature_size - 4) / 2 - 4) / 2)
self.conv1 = torch.nn.Conv2d(num_channels, 32, 5, 1)
self.conv2 = torch.nn.Conv2d(32, 64, 5, 1)
self.fc1 = torch.nn.Linear(self.features * self.features * 64, 1024)
self.lif0 = LIFFeedForwardCell(
(32, feature_size - 4, feature_size - 4),
p=LIFParameters(method=method, alpha=100.0),
)
self.lif1 = LIFFeedForwardCell(
(64, int((feature_size - 4) / 2) - 4, int(
(feature_size - 4) / 2) - 4),
p=LIFParameters(method=method, alpha=100.0),
)
self.lif2 = LIFFeedForwardCell((1024, ),
p=LIFParameters(method=method,
alpha=100.0))
self.out = LICell(1024, 10)
self.device = device
self.dtype = dtype
def __init__(self, cell, n_features=128, n_input=80, n_output=10):
super(SNNModel, self).__init__()
self.n_features = n_features
self.n_input = n_input
self.n_output = n_output
self.cell = cell(self.n_input, self.n_features)
self.readout = LICell(self.n_features, self.n_output)
def __init__(self,
device,
num_channels=1,
feature_size=28,
model="super",
dtype=torch.float):
super(ConvNet, self).__init__()
self.features = int(((feature_size - 4) / 2 - 4) / 2)
self.conv1 = torch.nn.Conv2d(num_channels, 20, 5, 1)
self.conv2 = torch.nn.Conv2d(20, 50, 5, 1)
self.fc1 = torch.nn.Linear(self.features * self.features * 50, 500)
self.out = LICell(500, 10)
self.device = device
self.lif0 = LIFFeedForwardCell(
(20, feature_size - 4, feature_size - 4),
p=LIFParameters(model=model, alpha=100.0),
)
self.lif1 = LIFFeedForwardCell(
(50, int((feature_size - 4) / 2) - 4, int(
(feature_size - 4) / 2) - 4),
p=LIFParameters(model=model, alpha=100.0),
)
self.lif2 = LIFFeedForwardCell((500, ),
p=LIFParameters(model=model,
alpha=100.0))
self.dtype = dtype
def test_li_cell():
cell = LICell(2, 4)
data = torch.randn(5, 2)
out, s = cell(data)
for x in s:
assert x.shape == (5, 4)
assert out.shape == (5, 4)
def test_li_cell_state():
cell = LICell(2, 4)
data = torch.randn(5, 2)
out, s = cell(data, LIState(torch.ones(5, 4), torch.ones(5, 4)))
for x in s:
assert x.shape == (5, 4)
assert out.shape == (5, 4)
def __init__(self):
super(Policy, self).__init__()
self.state_dim = 4
self.input_features = 16
self.hidden_features = 128
self.output_features = 2
self.constant_current_encoder = ConstantCurrentLIFEncoder(40)
self.lif = LIFCell(
2 * self.state_dim,
self.hidden_features,
p=LIFParameters(method="super", alpha=100.0),
)
self.dropout = torch.nn.Dropout(p=0.5)
self.readout = LICell(self.hidden_features, self.output_features)
self.saved_log_probs = []
self.rewards = []
def __init__(self, model="super"):
super(LSNNPolicy, self).__init__()
self.state_dim = 4
self.input_features = 16
self.hidden_features = 128
self.output_features = 2
# self.affine1 = torch.nn.Linear(self.state_dim, self.input_features)
self.constant_current_encoder = ConstantCurrentLIFEncoder(40)
self.lif_layer = LSNNCell(
2 * self.state_dim,
self.hidden_features,
p=LSNNParameters(method=model, alpha=100.0),
)
self.dropout = torch.nn.Dropout(p=0.5)
self.readout = LICell(self.hidden_features, self.output_features)
self.saved_log_probs = []
self.rewards = []
def __init__(
self, device, num_channels=1, feature_size=32, method="super", dtype=torch.float
):
super(ConvvNet4, self).__init__()
self.features = int(((feature_size - 4) / 2 - 4) / 2)
self.conv1 = torch.nn.Conv2d(1, 6, kernel_size=5, stride=1)
self.conv2 = torch.nn.Conv2d(6, 16, kernel_size=5,stride=1)
self.conv3 = torch.nn.Conv2d(16, 120, kernel_size=5, stride=1)
self.fc1 = torch.nn.Linear(120, 84)
# self.fc2 = torch.nn.Linear(84, 10)
self.lif0 = LIFFeedForwardCell(p=LIFParameters(method=method, alpha=100.0))
self.lif1 = LIFFeedForwardCell(p=LIFParameters(method=method, alpha=100.0))
self.lif2 = LIFFeedForwardCell(p=LIFParameters(method=method, alpha=100.0))
self.lif3 = LIFFeedForwardCell(p=LIFParameters(method=method, alpha=100.0))
self.out = LICell(84, 10)
self.device = device
self.dtype = dtype
def __init__(
self,
num_channels=1,
feature_size=32,
model="super",
dtype=torch.float,
):
super(Net, self).__init__()
self.features = int(((feature_size - 4) / 2 - 4) / 2)
self.conv1 = torch.nn.Conv2d(num_channels, 32, 5, 1)
self.conv2 = torch.nn.Conv2d(32, 64, 5, 1)
self.fc1 = torch.nn.Linear(self.features * self.features * 64, 1024)
self.lif0 = LIFFeedForwardCell(p=LIFParameters(method=model,
alpha=100.0), )
self.lif1 = LIFFeedForwardCell(p=LIFParameters(method=model,
alpha=100.0), )
self.lif2 = LIFFeedForwardCell(
p=LIFParameters(method=model, alpha=100.0))
self.out = LICell(1024, 10)
self.dtype = dtype
def __init__(
self,
input_features,
output_features,
seq_length,
is_lsnn,
dt=0.01,
model="super",
):
super(MemoryNet, self).__init__()
self.input_features = input_features
self.output_features = output_features
self.seq_length = seq_length
self.is_lsnn = is_lsnn
if is_lsnn:
p = LSNNParameters(method=model)
self.layer = LSNNCell(input_features, input_features, p, dt=dt)
else:
p = LIFParameters(method=model)
self.layer = LIFCell(input_features, input_features, dt=dt)
self.dropout = torch.nn.Dropout(p=0.2)
self.readout = LICell(input_features, output_features)
def test_cell_backward():
model = LICell(12, 1)
data = torch.ones(100, 12)
out, _ = model(data)
loss = out.sum()
loss.backward()
def main():
torch.manual_seed(42)
np.random.seed(42)
def generate_random_data(
seq_length,
batch_size,
input_features,
device="cpu",
dtype=torch.float,
dt=0.001,
):
freq = 5
prob = freq * dt
mask = torch.rand((seq_length, batch_size, input_features),
device=device,
dtype=dtype)
x_data = torch.zeros(
(seq_length, batch_size, input_features),
device=device,
dtype=dtype,
requires_grad=False,
)
x_data[mask < prob] = 1.0
y_data = torch.tensor(1 * (np.random.rand(batch_size) < 0.5),
device=device)
return x_data, y_data
seq_length = 500
batch_size = 1
input_features = 100
hidden_features = 8
output_features = 2
device = "cpu"
x, y_data = generate_random_data(
seq_length=seq_length,
batch_size=batch_size,
input_features=input_features,
device=device,
)
input_weights = (torch.randn((input_features, hidden_features),
device=device).float().t())
recurrent_weights = torch.randn((hidden_features, hidden_features),
device=device).float()
lif_correlation = LIFCorrelation(input_features, hidden_features)
out = LICell(hidden_features, output_features).to(device)
log_softmax_fn = torch.nn.LogSoftmax(dim=1)
loss_fn = torch.nn.NLLLoss()
linear_update = torch.nn.Linear(2 * 100 * 8, 100 * 8)
rec_linear_update = torch.nn.Linear(2 * 8 * 8, 8 * 8)
optimizer = torch.optim.Adam(
list(linear_update.parameters()) + [input_weights, recurrent_weights] +
list(out.parameters()),
lr=1e-1,
)
loss_hist = []
num_episodes = 100
for e in range(num_episodes):
s1 = lif_correlation.initial_state(batch_size, device=device)
so = out.initial_state(batch_size, device=device)
voltages = torch.zeros(seq_length,
batch_size,
output_features,
device=device)
hidden_voltages = torch.zeros(seq_length,
batch_size,
hidden_features,
device=device)
hidden_currents = torch.zeros(seq_length,
batch_size,
hidden_features,
device=device)
optimizer.zero_grad()
for ts in range(seq_length):
z1, s1 = lif_correlation(
x[ts, :, :],
s1,
input_weights=input_weights,
recurrent_weights=recurrent_weights,
)
input_weights = correlation_based_update(
ts,
linear_update,
input_weights.detach(),
s1.input_correlation_state,
0.01,
10,
)
recurrent_weights = correlation_based_update(
ts,
rec_linear_update,
recurrent_weights.detach(),
s1.recurrent_correlation_state,
0.01,
10,
)
vo, so = out(z1, so)
hidden_voltages[ts, :, :] = s1.lif_state.v.detach()
hidden_currents[ts, :, :] = s1.lif_state.i.detach()
voltages[ts, :, :] = vo
m, _ = torch.max(voltages, dim=0)
log_p_y = log_softmax_fn(m)
loss_val = loss_fn(log_p_y, y_data)
loss_val.backward()
optimizer.step()
loss_hist.append(loss_val.item())
print(f"{e}/{num_episodes}: {loss_val.item()}")
np.save("loss.npy", loss_hist)
import matplotlib.pyplot as plt
plt.semilogy(loss_hist)
plt.savefig("loss.png")