class ConvNet(torch.nn.Module):
def __init__(self,
device,
num_channels=1,
feature_size=28,
method="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(method=method, alpha=100.0),
)
self.lif1 = LIFFeedForwardCell(
(50, int((feature_size - 4) / 2) - 4, int(
(feature_size - 4) / 2) - 4),
p=LIFParameters(method=method, alpha=100.0),
)
self.lif2 = LIFFeedForwardCell((500, ),
p=LIFParameters(method=method,
alpha=100.0))
self.dtype = dtype
def forward(self, x):
seq_length = x.shape[0]
batch_size = x.shape[1]
# specify the initial states
s0 = self.lif0.initial_state(batch_size, self.device, self.dtype)
s1 = self.lif1.initial_state(batch_size, self.device, self.dtype)
s2 = self.lif2.initial_state(batch_size, self.device, self.dtype)
so = self.out.initial_state(batch_size,
device=self.device,
dtype=self.dtype)
voltages = torch.zeros(seq_length,
batch_size,
10,
device=self.device,
dtype=self.dtype)
for ts in range(seq_length):
z = self.conv1(x[ts, :])
z, s0 = self.lif0(z, s0)
z = torch.nn.functional.max_pool2d(z, 2, 2)
z = 10 * self.conv2(z)
z, s1 = self.lif1(z, s1)
z = torch.nn.functional.max_pool2d(z, 2, 2)
z = z.view(-1, self.features**2 * 50)
z = self.fc1(z)
z, s2 = self.lif2(z, s2)
v, so = self.out(torch.nn.functional.relu(z), so)
voltages[ts, :, :] = v
return voltages
def test_lif_feedforward_cell_backward():
# Tests that gradient variables can be used in subsequent applications
cell = LIFFeedForwardCell()
data = torch.randn(5, 4)
out, s = cell(data)
out, _ = cell(out, s)
loss = out.sum()
loss.backward()
def test_lif_feedforward_cell():
layer = LIFFeedForwardCell()
data = torch.randn(5, 4)
out, s = layer(data)
assert out.shape == (5, 4)
for x in s:
assert x.shape == (5, 4)
def __init__(self,
num_channels=1,
feature_size=28,
method="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.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.dtype = dtype
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,
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