Truncated BPTT

[tmasquelier]

Hi there,
SpikingJelly is great, congrats to the developers, and thank you :-)
When dealing with a long DVS sequence with labels provided at multiple time points, truncated backprop through time (TBPTT) can be useful, see for example: http://arxiv.org/abs/2009.13436
I think it would be straightforward to use TBPTT in SpikingJelly, but only in the step-by-step mode, not in the layer-by-layer one, right?

[fangwei123456]

Thanks! Let me try to add a code sample in this discussion.

[fangwei123456]

Here is an example of TBPTT based on classifying FashionMNIST:

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from spikingjelly.clock_driven import neuron, functional, surrogate, layer
from torch.utils.tensorboard import SummaryWriter
import sys
import time
if sys.platform != 'win32':
    import readline
import numpy as np
torch.backends.cudnn.benchmark = True
_seed_ = 2020
torch.manual_seed(_seed_)  # use torch.manual_seed() to seed the RNG for all devices (both CPU and CUDA)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(_seed_)
class Net(nn.Module):
    def __init__(self, tau, T, v_threshold=1.0, v_reset=0.0, T_truncated=0):
        super().__init__()
        self.T = T
        self.T_truncated = T_truncated

        self.static_conv = nn.Sequential(
            nn.Conv2d(1, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128),
        )

        self.conv = nn.Sequential(
            neuron.IFNode(v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
            nn.MaxPool2d(2, 2),  # 14 * 14

            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128),
            neuron.IFNode(v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
            nn.MaxPool2d(2, 2)  # 7 * 7

        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            layer.Dropout(0.5),
            nn.Linear(128 * 7 * 7, 128 * 3 * 3, bias=False),
            neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
            layer.Dropout(0.5),
            nn.Linear(128 * 3 * 3, 128, bias=False),
            neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
            nn.Linear(128, 10, bias=False),
            neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
        )

    def detach(self):
        for m in self.modules():
            if isinstance(m, neuron.BaseNode):
                m.v.detach_()
            elif isinstance(m, layer.Dropout):
                m.mask.detach_()


    def forward(self, x):
        x = self.static_conv(x)

        if self.training:
            bptt_length = 0
            out_spikes_counter = self.fc(self.conv(x))
            bptt_length += 1
            if bptt_length == self.T_truncated:
                self.detach()
                bptt_length = 0

            for t in range(1, self.T):
                out_spikes_counter += self.fc(self.conv(x))
                bptt_length += 1
                if bptt_length == self.T_truncated:
                    self.detach()
                    bptt_length = 0
        else:
            out_spikes_counter = self.fc(self.conv(x))
            for t in range(1, self.T):
                out_spikes_counter += self.fc(self.conv(x))

        return out_spikes_counter / self.T
def main():

    device = 'cuda:0'
    dataset_dir = '/userhome/datasets/FashionMNIST'
    batch_size = 128
    learning_rate = 1e-3
    T = 8
    tau = 2.0
    train_epoch = 100
    log_dir = './logs_tbptt'
    T_truncated = 2


    writer = SummaryWriter(log_dir)

    # 初始化数据加载器
    train_data_loader = torch.utils.data.DataLoader(
        dataset=torchvision.datasets.FashionMNIST(
            root=dataset_dir,
            train=True,
            transform=torchvision.transforms.ToTensor(),
            download=True),
        batch_size=batch_size,
        shuffle=True,
        drop_last=True)
    test_data_loader = torch.utils.data.DataLoader(
        dataset=torchvision.datasets.FashionMNIST(
            root=dataset_dir,
            train=False,
            transform=torchvision.transforms.ToTensor(),
            download=True),
        batch_size=batch_size,
        shuffle=True,
        drop_last=False)

    # 初始化网络
    net = Net(tau=tau, T=T, T_truncated=T_truncated).to(device)
    # 使用Adam优化器
    optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
    train_times = 0
    max_test_accuracy = 0
    for epoch in range(train_epoch):
        net.train()
        t_start = time.perf_counter()
        for img, label in train_data_loader:
            img = img.to(device)
            label = label.to(device)
            label_one_hot = F.one_hot(label, 10).float()

            optimizer.zero_grad()

            out_spikes_counter_frequency = net(img)

            # 损失函数为输出层神经元的脉冲发放频率，与真实类别的MSE
            # 这样的损失函数会使，当类别i输入时，输出层中第i个神经元的脉冲发放频率趋近1，而其他神经元的脉冲发放频率趋近0
            loss = F.mse_loss(out_spikes_counter_frequency, label_one_hot)
            loss.backward()
            optimizer.step()
            # 优化一次参数后，需要重置网络的状态，因为SNN的神经元是有“记忆”的
            functional.reset_net(net)

            # 正确率的计算方法如下。认为输出层中脉冲发放频率最大的神经元的下标i是分类结果
            accuracy = (out_spikes_counter_frequency.max(1)[1] == label.to(device)).float().mean().item()
            if train_times % 256 == 0:
                writer.add_scalar('train_accuracy', accuracy, train_times)
            train_times += 1
        t_train = time.perf_counter() - t_start
        net.eval()
        t_start = time.perf_counter()
        with torch.no_grad():
            # 每遍历一次全部数据集，就在测试集上测试一次
            test_sum = 0
            correct_sum = 0
            for img, label in test_data_loader:
                img = img.to(device)
                out_spikes_counter_frequency = net(img)

                correct_sum += (out_spikes_counter_frequency.max(1)[1] == label.to(device)).float().sum().item()
                test_sum += label.numel()
                functional.reset_net(net)
            test_accuracy = correct_sum / test_sum
            t_test = time.perf_counter() - t_start
            writer.add_scalar('test_accuracy', test_accuracy, epoch)
            if max_test_accuracy < test_accuracy:
                max_test_accuracy = test_accuracy
                print('saving net...')
                torch.save(net, log_dir + '/net_max_acc.pt')
                print('saved')

        print('epoch={}, t_train={}, t_test={}, device={}, dataset_dir={}, batch_size={}, learning_rate={}, T={}, log_dir={}, max_test_accuracy={}, train_times={}'.format(epoch, t_train, t_test, device, dataset_dir, batch_size, learning_rate, T, log_dir, max_test_accuracy, train_times))

if __name__ == '__main__':
    main()

[fangwei123456]

Yes, I find this code are not the really TBPTT. But this code works similarly:

forward 2 steps
backward 2 steps

Let me make a little change.

[fangwei123456]

This code is same with the last code in computation but can save more memory.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from spikingjelly.clock_driven import neuron, functional, surrogate, layer
from torch.utils.tensorboard import SummaryWriter
import sys
import time
if sys.platform != 'win32':
    import readline
import numpy as np
_seed_ = 2020
torch.manual_seed(_seed_)  # use torch.manual_seed() to seed the RNG for all devices (both CPU and CUDA)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(_seed_)
class Net(nn.Module):
    def __init__(self, tau, v_threshold=1.0, v_reset=0.0):
        super().__init__()
        self.x = None

        self.static_conv = nn.Sequential(
            nn.Conv2d(1, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128),
        )

        self.conv = nn.Sequential(
            neuron.IFNode(v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
            nn.MaxPool2d(2, 2),  # 14 * 14

            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128),
            neuron.IFNode(v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
            nn.MaxPool2d(2, 2)  # 7 * 7

        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            layer.Dropout(0.5),
            nn.Linear(128 * 7 * 7, 128 * 3 * 3, bias=False),
            neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
            layer.Dropout(0.5),
            nn.Linear(128 * 3 * 3, 128, bias=False),
            neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
            nn.Linear(128, 10, bias=False),
            neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset, surrogate_function=surrogate.ATan()),
        )

    def detach(self):
        for m in self.modules():
            if isinstance(m, neuron.BaseNode):
                m.v.detach_()
            elif isinstance(m, layer.Dropout):
                m.mask.detach_()


    def forward(self, x, T):
        if self.x is None:
            x = self.static_conv(x)
        else:
            x = self.x

        out_spikes_counter = self.fc(self.conv(x))
        for t in range(1, T):
            out_spikes_counter += self.fc(self.conv(x))

        return out_spikes_counter / T

    def reset(self):
        self.x = None

def main():

    device = 'cuda:0'
    dataset_dir = '/userhome/datasets/FashionMNIST'
    batch_size = 128
    learning_rate = 1e-3
    T = 8
    tau = 2.0
    train_epoch = 100
    log_dir = './logs_tbptt'
    T_truncated = 2


    writer = SummaryWriter(log_dir)

    # 初始化数据加载器
    train_data_loader = torch.utils.data.DataLoader(
        dataset=torchvision.datasets.FashionMNIST(
            root=dataset_dir,
            train=True,
            transform=torchvision.transforms.ToTensor(),
            download=True),
        batch_size=batch_size,
        shuffle=True,
        drop_last=True)
    test_data_loader = torch.utils.data.DataLoader(
        dataset=torchvision.datasets.FashionMNIST(
            root=dataset_dir,
            train=False,
            transform=torchvision.transforms.ToTensor(),
            download=True),
        batch_size=batch_size,
        shuffle=True,
        drop_last=False)

    # 初始化网络
    net = Net(tau=tau).to(device)
    # 使用Adam优化器
    optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
    train_times = 0
    max_test_accuracy = 0
    for epoch in range(train_epoch):
        net.train()
        t_start = time.perf_counter()
        for img, label in train_data_loader:
            img = img.to(device)
            label = label.to(device)
            label_one_hot = F.one_hot(label, 10).float()

            for t in range(0, T, T_truncated):
                optimizer.zero_grad()
                out_spikes_counter_frequency = net(img, T_truncated)
                loss = F.mse_loss(out_spikes_counter_frequency, label_one_hot)
                loss.backward()
                optimizer.step()
                net.detach()

            functional.reset_net(net)

            # 正确率的计算方法如下。认为输出层中脉冲发放频率最大的神经元的下标i是分类结果
            accuracy = (out_spikes_counter_frequency.max(1)[1] == label.to(device)).float().mean().item()
            if train_times % 256 == 0:
                writer.add_scalar('train_accuracy', accuracy, train_times)
            train_times += 1
        t_train = time.perf_counter() - t_start
        net.eval()
        t_start = time.perf_counter()
        with torch.no_grad():
            # 每遍历一次全部数据集，就在测试集上测试一次
            test_sum = 0
            correct_sum = 0
            for img, label in test_data_loader:
                img = img.to(device)
                out_spikes_counter_frequency = net(img, T)

                correct_sum += (out_spikes_counter_frequency.max(1)[1] == label.to(device)).float().sum().item()
                test_sum += label.numel()
                functional.reset_net(net)
            test_accuracy = correct_sum / test_sum
            t_test = time.perf_counter() - t_start
            writer.add_scalar('test_accuracy', test_accuracy, epoch)
            if max_test_accuracy < test_accuracy:
                max_test_accuracy = test_accuracy
                print('saving net...')
                torch.save(net, log_dir + '/net_max_acc.pt')
                print('saved')

        print('epoch={}, t_train={}, t_test={}, device={}, dataset_dir={}, batch_size={}, learning_rate={}, T={}, log_dir={}, max_test_accuracy={}, train_times={}'.format(epoch, t_train, t_test, device, dataset_dir, batch_size, learning_rate, T, log_dir, max_test_accuracy, train_times))

if __name__ == '__main__':
    main()

[fangwei123456]

It will be difficult to implement TBPTT with arbitrary k1 and k2.. In this code, k1=k2=2.

[tmasquelier]

I think the case k1=k2 case is solved :-) Thanks. Plus I think this approach is compatible with the multi-step mode (and CUDA neurons).
Do you think k2=2*k1 is doable?

[fangwei123456]

OK. I will try to add a code example

[fangwei123456]

Here is an example of k1 = T_truncated, k2 = 2 * T_truncated. But it is not the standard TBPTT because some computation graph is not re-used. The re-used computation graph is difficult to implement.

            retain_graph = True
            for t in range(0, T, T_truncated):
                optimizer.zero_grad()
                out_spikes_counter_frequency = net(img, T_truncated)
                loss = F.mse_loss(out_spikes_counter_frequency, label_one_hot)
                loss.backward(retain_graph=retain_graph)
                optimizer.step()
                if not retain_graph:
                    net.detach()
                retain_graph = not retain_graph

[urancon]

Hi, very interesting discussion, I have been carefully following it so far. Thank you for all your insights !
When you say that

some computation graph is not re-used

can you clarify please ?

Indeed in the method you propose, I think that the number of timesteps on which each backward pass is done alternates between 2 * k1 and k1. I mean, the first backward pass will be done with the error accumulated during 2 * k1 timesteps, the second with the error accumulated during k1 timesteps, the third with 2 * k1 again, etc.

Is that what you meant, or is it more complicated than that ?

[fangwei123456]

Thanks for you following! This code runs almost similar to your description, but with tiny difference. Let me explain:

1. forward T_truncated steps (build a computation graph with T_truncated steps)
2. backward T_truncated steps and keep the computation graph
3. do the gradient descend and zero the gradient (so, the error will not be accumulated)
4. forward T_truncated steps (add the computation graph with T_truncated steps, now the graph have 2*T_truncated steps)
5. backward 2*T_truncated steps and destroy the computation graph
6. do the gradient descend and zero the gradient
7. detach all hidden states
8. retuen to 1

[fangwei123456]

The re-use problem:

1.If k1 > k2,
Forward k1 steps, backward k2 steps, which can be done by detaching hidden states at t = (k1 - k2 - 1)

2.If k1 < k2,
Forward k1 steps, but if we want to backward k2 steps, we have to firstly et a graph with (k2 - k1) steps. To reuse such a graph, we need to do more changes on the current codes.

[fangwei123456]

"some computation graph is not re-used"

In the standard TBPTT,

forward step: [0, T_truncated), [T_truncated, 2*T_truncated), [2*T_truncated, 3*T_truncated), [3*T_truncated, 4*T_truncated)...
backward step: [0, T_truncated), [0, 2*T_truncated), [T_truncated, 3*T_truncated), [2*T_truncated, 4*T_truncated)...

The current code is:

forward step: [0, T_truncated), [T_truncated, 2*T_truncated), [2*T_truncated, 3*T_truncated), [3*T_truncated, 4*T_truncated)...
backward step: [0, T_truncated), [0, 2*T_truncated), [2*T_truncated, 3*T_truncated), [2*T_truncated, 4*T_truncated)...

[tmasquelier]

Thanks Wei, very clear.
It is unfortunate that 50% of the time the backward is done one T_truncated steps only, not on 2*T_truncated step.
I'm going to take a look on how TBPTT is implemented in RNNs in general. This pb is not specific to SNNs. For example TBPTT can be used to train LSTMs.

[tmasquelier]

I've just heard that PyTorch Lightning has TBPTT:
https://pytorch-lightning.readthedocs.io/en/stable/guides/data.html?highlight=TBPTT#truncated-backpropagation-through-time-tbptt