SNN的一个简单示例_snntorch

# imports
import snntorch as snn
from snntorch import surrogate
# from snntorch import backprop
from snntorch import functional as SF
from snntorch import utils
# from snntorch import spikeplot as splt

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# import torch.nn.functional as F

import matplotlib.pyplot as plt

# import numpy as np
# import itertools

# dataloader arguments
batch_size = 64  # 内存不够128
data_path = 'data'

dtype = torch.float
# device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
device = torch.device("cpu")  # 这个电脑的cuda版本太低了

# Define a transform
transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.Grayscale(),
    transforms.ToTensor(),
    transforms.Normalize((0,), (1,))])  # 其实这行都没什么用

mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)
mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)

# Create DataLoaders
train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)
test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)

# neuron and simulation parameters
spike_grad = surrogate.fast_sigmoid(slope=25)  # 替代梯度，用于反向传播
beta = 0.5  # 神经元膜电位的衰减率
num_steps = 50  # 时间步（SNN特有的）

# Define Network
# class Net(nn.Module):
#     def __init__(self):
#         super().__init__()
#
#         # Initialize layers
#         self.conv1 = nn.Conv2d(1, 12, 5)
#         self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)
#         self.conv2 = nn.Conv2d(12, 64, 5)
#         self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)
#         self.fc1 = nn.Linear(64*4*4, 10)
#         self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)
#
#     def forward(self, x):
#
#         # Initialize hidden states and outputs at t=0
#         mem1 = self.lif1.init_leaky()
#         mem2 = self.lif2.init_leaky()
#         mem3 = self.lif3.init_leaky()
#
#         cur1 = F.max_pool2d(self.conv1(x), 2)
#         spk1, mem1 = self.lif1(cur1, mem1)
#
#         cur2 = F.max_pool2d(self.conv2(spk1), 2)
#         spk2, mem2 = self.lif2(cur2, mem2)
#
#         cur3 = self.fc1(spk2.view(batch_size, -1))
#         spk3, mem3 = self.lif3(cur3, mem3)
#
#         return spk3, mem3

#  Initialize Network
net = nn.Sequential(nn.Conv2d(1, 12, 5),
                    nn.MaxPool2d(2),
                    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),
                    nn.Conv2d(12, 64, 5),
                    nn.MaxPool2d(2),
                    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True),
                    nn.Flatten(),
                    nn.Linear(64 * 4 * 4, 10),
                    snn.Leaky(beta=beta, spike_grad=spike_grad, init_hidden=True, output=True)
                    ).to(device)


def forward_pass(net, num_steps, data):
    mem_rec = []
    spk_rec = []
    utils.reset(net)  # resets hidden states for all LIF neurons in net

    for step in range(num_steps):
        spk_out, mem_out = net(data)
        spk_rec.append(spk_out)
        mem_rec.append(mem_out)
    # 难道说每个step都不会初始化了，我觉得应该是这样
    return torch.stack(spk_rec), torch.stack(mem_rec)


loss_fn = SF.ce_rate_loss()


def batch_accuracy(train_loader, net, num_steps):
    # 训练集一个batch的accuracy
    with torch.no_grad():
        total = 0
        acc = 0
        net.eval()

        train_loader = iter(train_loader)
        for data, targets in train_loader:
            data = data.to(device)
            targets = targets.to(device)
            spk_rec, _ = forward_pass(net, num_steps, data)

            acc += SF.accuracy_rate(spk_rec, targets) * spk_rec.size(1)
            total += spk_rec.size(1)

    return acc / total


optimizer = torch.optim.Adam(net.parameters(), lr=1e-2, betas=(0.9, 0.999))
num_epochs = 1
loss_hist = []
test_acc_hist = []
counter = 0

# Outer training loop
for epoch in range(num_epochs):  # 只训练1个epoch！？

    # Training loop
    for data, targets in iter(train_loader):
        data = data.to(device)
        targets = targets.to(device)

        # forward pass
        net.train()
        #进行num_steps次循环，完成一次前向传播
        spk_rec, _ = forward_pass(net, num_steps, data)

        # initialize the loss & sum over time
        loss_val = loss_fn(spk_rec, targets) #和正确的标签进行计算得到损失

        # Gradient calculation + weight update
        optimizer.zero_grad()
        #根据损失计算梯度
        loss_val.backward()
        #更新梯度
        optimizer.step()

        # Store loss history for future plotting
        loss_hist.append(loss_val.item())

        # Test set
        if counter % 50 == 0:
            with torch.no_grad():
                net.eval()
                
                # Test set forward pass
                test_acc = batch_accuracy(test_loader, net, num_steps)
                print(f"Iteration {counter}, Test Acc: {test_acc * 100:.2f}%\n")
                test_acc_hist.append(test_acc.item())

        counter += 1

# Plot Loss
fig = plt.figure(facecolor="w")
plt.plot(test_acc_hist)
plt.title("Test Set Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.show()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177

1