def make_layers(cfg, batch_norm=False):
layers = []
in_channels = 3
for v in cfg:
if v == "M":
layers += [Lift(nn.MaxPool2d(kernel_size=2, stride=2))]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
if batch_norm:
layers += [Lift(conv2d), Lift(nn.BatchNorm2d(v)), LIF()]
else:
layers += [Lift(conv2d), LIF()]
in_channels = v
return nn.Sequential(*layers)
def __init__(self, features, num_classes=1000, init_weights=True):
super(VGG, self).__init__()
self.features = features
self.avgpool = Lift(nn.AdaptiveAvgPool2d((7, 7)))
self.classifier = nn.Sequential(
Lift(nn.Linear(512 * 7 * 7, 4096)),
LIF(),
Lift(nn.Dropout()),
Lift(nn.Linear(4096, 4096)),
LIF(),
Lift(nn.Dropout()),
Lift(nn.Linear(4096, num_classes)),
)
if init_weights:
self._initialize_weights()
def test_lif_feedforward_layer():
layer = LIF()
data = torch.randn(10, 5, 4)
out, s = layer(data)
assert out.shape == (10, 5, 4)
for x in s:
assert x.shape == (5, 4)
def test_lif_feedforward_layer_backward_iteration():
# Tests that gradient variables can be used in subsequent applications
model = LIF()
data = torch.ones(10, 6)
out, s = model(data)
out, _ = model(out, s)
loss = out.sum()
loss.backward()
def test_lift_sequential():
batch_size = 16
seq_length = 20
in_channels = 64
out_channels = 32
data = torch.randn(seq_length, batch_size, in_channels, 20, 30)
module = torch.nn.Sequential(
Lift(torch.nn.Conv2d(in_channels, out_channels, 5, 1)),
LIF(),
)
output, _ = module(data)
assert output.shape == torch.Size(
[seq_length, batch_size, out_channels, 16, 26])
def __init__(self,
in_planes,
out_planes,
norm_layer,
kernel_size=3,
stride=1,
groups=1):
padding = (kernel_size - 1) // 2
super(ConvBNReLU, self).__init__(
Lift(
nn.Conv2d(
in_planes,
out_planes,
kernel_size,
stride,
padding,
groups=groups,
bias=False,
)),
Lift(norm_layer(out_planes)),
LIF(),
)
def test_lif_in_time():
layer = LIF()
data = torch.randn(10, 5, 2)
out, _ = layer(data)
assert out.shape == (10, 5, 2)
def test_lif_feedforward_layer_backward():
model = LIF()
data = torch.ones(10, 12)
out, _ = model(data)
loss = out.sum()
loss.backward()
def __init__(self):
super(SNNetwork, self).__init__()
self.l0 = LIF()
self.l1 = LIF()
self.s0 = self.s1 = None