snn_conv_1_gesture_max.py

import argparse
import pandas as pd
import os
import time
import sys
import csv
import struct

import torch
import numpy as np
import random
from torch.utils.data import Dataset, DataLoader, random_split
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets
from torchvision import utils

import spikingjelly
import spikingjelly.datasets

from snn_lib.snn_layers import *
from snn_lib.optimizers import *
from snn_lib.schedulers import *
from snn_lib.data_loaders import *
import snn_lib.utilities

import omegaconf
from omegaconf import OmegaConf

if torch.cuda.is_available():
    device = torch.device('cuda:2')
else:
    device = torch.device('cpu')

# arg parser
parser = argparse.ArgumentParser(description='conv snn')
parser.add_argument('--config_file', type=str, default='snn_conv_1_gesture_max.yaml',
                    help='path to configuration file')
parser.add_argument('--train', action='store_true',
                    help='train model')

parser.add_argument('--test', action='store_true',
                    help='test model')
parser.add_argument('--load', action='store_true', help='load dataloader')

args = parser.parse_args()

# %% config file
if args.config_file is None:
    print('No config file provided, use default config file')
else:
    print('Config file provided:', args.config_file)

conf = OmegaConf.load(args.config_file)

torch.manual_seed(conf['pytorch_seed'])
np.random.seed(conf['pytorch_seed'])

experiment_name = conf['experiment_name']

# %% checkpoint
save_checkpoint = conf['save_checkpoint']
checkpoint_base_name = conf['checkpoint_base_name']
checkpoint_base_path = conf['checkpoint_base_path']
test_checkpoint_path = conf['test_checkpoint_path']

# %% training parameters
hyperparam_conf = conf['hyperparameters']
length = hyperparam_conf['length']
batch_size = hyperparam_conf['batch_size']
synapse_type = hyperparam_conf['synapse_type']
epoch = hyperparam_conf['epoch']
tau_m = hyperparam_conf['tau_m']
tau_s = hyperparam_conf['tau_s']
filter_tau_m = hyperparam_conf['filter_tau_m']
filter_tau_s = hyperparam_conf['filter_tau_s']

membrane_filter = hyperparam_conf['membrane_filter']

train_bias = hyperparam_conf['train_bias']
train_coefficients = hyperparam_conf['train_coefficients']

# acc file name
acc_file_name = experiment_name + '_' + conf['acc_file_name']


class GestureDataset(Dataset):
    def __init__(self, root, train, length):
        super(GestureDataset, self).__init__()
        self.data, self.label = self.get_dataset(root=root, train=train, length=length)

    def __len__(self):
        return len(self.label)

    def __getitem__(self, idx):
        return self.data[idx], self.label[idx]

    def get_spike_train(self, event, length):
        p, x, y, t = event['p'], event['x'], event['y'], event['t']
        t -= t.min()
        t = t // 1000
        bin_width = (t.max() + 1) // length
        spike_train = torch.zeros((2, 128, 128,  (length + 1) * bin_width), dtype=torch.bool)  # [p, x, y, t]
        spike_train[p, x, y, t] = True
        spike_train = spike_train.reshape(2, 128, 128, length + 1, bin_width)
        spike_train = spike_train.sum(dim=4, dtype=torch.float32)
        spike_train = spike_train[:, :, :, 0:length]
        spike_train /= spike_train.amax(dim=(0, 1, 2))  # max
        return spike_train  # [p, x, y, t]

    def get_dataset(self, root, train, length):
        data = []
        label = []
        dataset = spikingjelly.datasets.dvs128_gesture.DVS128Gesture(root=root, train=train, use_frame=False)
        for event, number in dataset:
            data.append(self.get_spike_train(event=event, length=length))
            label.append(number)
        data = torch.stack(data)
        label = torch.tensor(label)
        return data, label


# %% define model
class mysnn(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.length = length
        self.batch_size = batch_size

        self.train_coefficients = train_coefficients
        self.train_bias = train_bias
        self.membrane_filter = membrane_filter

        # 1: 2x128x128 -> 64x42x42
        self.axon1 = dual_exp_iir_layer((2, 128, 128), self.length, self.batch_size, tau_m, tau_s, train_coefficients)
        self.conv1 = conv2d_layer(
            h_input=128, w_input=128, in_channels=2, out_channels=64, kernel_size=7,
            stride=3, padding=1, dilation=1, step_num=length, batch_size=batch_size,
            tau_m=tau_m, train_bias=train_bias, membrane_filter=membrane_filter, input_type='axon'
        )
        # 2: 64x42x42 -> 32x14x14
        self.axon2 = dual_exp_iir_layer((64, 42, 42), self.length, self.batch_size, tau_m, tau_s, train_coefficients)
        self.conv2 = conv2d_layer(
            h_input=42, w_input=42, in_channels=64, out_channels=32, kernel_size=3,
            stride=3, padding=0, dilation=1, step_num=length, batch_size=batch_size,
            tau_m=tau_m, train_bias=train_bias, membrane_filter=membrane_filter, input_type='axon'
        )
        # 3: 32x14x14-> 32x14x14
        self.axon3 = dual_exp_iir_layer((32, 14, 14), self.length, self.batch_size, tau_m, tau_s, train_coefficients)
        self.conv3 = conv2d_layer(
            h_input=14, w_input=14, in_channels=32, out_channels=32, kernel_size=3,
            stride=1, padding=1, dilation=1, step_num=length, batch_size=batch_size,
            tau_m=tau_m, train_bias=train_bias, membrane_filter=membrane_filter, input_type='axon'
        )
        # 4
        self.axon4 = dual_exp_iir_layer((32 * 14 * 14,), self.length, self.batch_size, tau_m, tau_s, train_coefficients)
        self.snn4 = neuron_layer(32 * 14 * 14, 256, self.length, self.batch_size, tau_m, self.train_bias,
                                 self.membrane_filter)
        self.dropout4 = torch.nn.Dropout(p=0.3, inplace=False)
        # 5
        self.axon5 = dual_exp_iir_layer((256,), self.length, self.batch_size, tau_m, tau_s, train_coefficients)
        self.snn5 = neuron_layer(256, 11, self.length, self.batch_size, tau_m, self.train_bias, self.membrane_filter)

    def forward(self, inputs):
        """
        :param inputs: [batch, 2, 34, 34, t]
        :return:
        """
        # 1
        axon1_out, _ = self.axon1(inputs, self.axon1.create_init_states())
        spike_l1, _ = self.conv1(axon1_out,self.conv1.create_init_states())
        # 2
        axon2_out, _ = self.axon2(spike_l1, self.axon2.create_init_states())
        spike_l2, _ = self.conv2(axon2_out, self.conv2.create_init_states())
        # 3
        axon3_out, _ = self.axon3(spike_l2, self.axon3.create_init_states())
        spike_l3, _ = self.conv3(axon3_out, self.conv3.create_init_states())
        # 3 -> 4
        spike_l3 = spike_l3.view(spike_l3.shape[0], -1, spike_l3.shape[-1])
        # 4
        axon4_out, _ = self.axon4(spike_l3, self.axon4.create_init_states())
        spike_l4, _ = self.snn4(axon4_out, self.snn4.create_init_states())
        drop_4 = self.dropout4(spike_l4)
        # 5
        axon5_out, _ = self.axon5(drop_4, self.axon5.create_init_states())
        spike_l5, _ = self.snn5(axon5_out, self.snn5.create_init_states())

        return spike_l5


########################### train function ###################################
def train(model, optimizer, scheduler, train_data_loader, writer=None):
    eval_image_number = 0
    correct_total = 0
    wrong_total = 0

    criterion = torch.nn.CrossEntropyLoss()

    model.train()

    for i_batch, sample_batched in enumerate(train_data_loader):

        x_train = sample_batched[0].to(device)
        target = sample_batched[1].to(device)
        out_spike = model(x_train)

        spike_count = torch.sum(out_spike, dim=2)

        model.zero_grad()
        loss = criterion(spike_count, target.long())
        loss.backward()
        optimizer.step()

        # calculate acc
        _, idx = torch.max(spike_count, dim=1)

        eval_image_number += len(sample_batched[1])
        wrong = len(torch.where(idx != target)[0])

        correct = len(sample_batched[1]) - wrong
        wrong_total += len(torch.where(idx != target)[0])
        correct_total += correct
        acc = correct_total / eval_image_number

        # scheduler step
        if isinstance(scheduler, torch.optim.lr_scheduler.CyclicLR):
            scheduler.step()

    # scheduler step
    if isinstance(scheduler, torch.optim.lr_scheduler.MultiStepLR):
        scheduler.step()

    acc = correct_total / eval_image_number

    return acc, loss


def test(model, test_data_loader, writer=None):
    eval_image_number = 0
    correct_total = 0
    wrong_total = 0

    model.eval()

    criterion = torch.nn.CrossEntropyLoss()

    for i_batch, sample_batched in enumerate(test_data_loader):
        x_test = sample_batched[0].to(device)
        target = sample_batched[1].to(device)
        out_spike = model(x_test)

        spike_count = torch.sum(out_spike, dim=2)

        loss = criterion(spike_count, target.long())

        # calculate acc
        _, idx = torch.max(spike_count, dim=1)

        eval_image_number += len(sample_batched[1])
        wrong = len(torch.where(idx != target)[0])

        correct = len(sample_batched[1]) - wrong
        wrong_total += len(torch.where(idx != target)[0])
        correct_total += correct
        acc = correct_total / eval_image_number

    acc = correct_total / eval_image_number

    return acc, loss


if __name__ == "__main__":

    snn = mysnn().to(device)
    snn = torch.nn.DataParallel(snn, device_ids=[2, 3])

    writer = SummaryWriter()

    params = list(snn.parameters())

    optimizer = get_optimizer(params, conf)

    scheduler = get_scheduler(optimizer, conf)

    train_acc_list = []
    test_acc_list = []
    checkpoint_list = []

    if args.train == True:
        if args.load is True:
            print('load data')
            train_data = torch.load('./data/DvsGesture/train_data_max.pt')
        else:
            print('process data')
            train_data = GestureDataset(root='./data/DvsGesture', train=True, length=length)
            torch.save(train_data, './data/DvsGesture/train_data_max.pt')
        print('train data', len(train_data))
        train_data, dev_data = random_split(
            train_data, [928, 128], generator=torch.Generator().manual_seed(42)
        )
        print('train_data', len(train_data), 'dev_data', len(dev_data))
        train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, drop_last=False)
        dev_dataloader = DataLoader(dev_data, batch_size=batch_size, shuffle=False, drop_last=False)

        train_it = 0
        test_it = 0
        for j in range(epoch):

            epoch_time_stamp = time.strftime("%Y%m%d-%H%M%S")

            snn.train()
            train_acc, train_loss = train(snn, optimizer, scheduler, train_dataloader, writer=None)
            train_acc_list.append(train_acc)

            print('Train epoch: {}, acc: {}'.format(j, train_acc))

            # save every checkpoint
            if save_checkpoint == True:
                checkpoint_name = checkpoint_base_name + experiment_name + '_' + str(j) + '_' + epoch_time_stamp
                checkpoint_path = os.path.join(checkpoint_base_path, checkpoint_name)
                checkpoint_list.append(checkpoint_path)

                torch.save({
                    'epoch': j,
                    'snn_state_dict': snn.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'loss': train_loss,
                }, checkpoint_path)

            # test model
            snn.eval()
            test_acc, test_loss = test(snn, dev_dataloader, writer=None)

            print('Test epoch: {}, acc: {}'.format(j, test_acc))
            test_acc_list.append(test_acc)

        # save result and get best epoch
        train_acc_list = np.array(train_acc_list)
        test_acc_list = np.array(test_acc_list)

        acc_df = pd.DataFrame(data={'train_acc': train_acc_list, 'test_acc': test_acc_list})

        acc_df.to_csv(acc_file_name)

        best_train_acc = np.max(train_acc_list)
        best_train_epoch = np.argmax(test_acc_list)

        best_test_epoch = np.argmax(test_acc_list)
        best_test_acc = np.max(test_acc_list)

        best_checkpoint = checkpoint_list[best_test_epoch]

        print('Summary:')
        print('Best train acc: {}, epoch: {}'.format(best_train_acc, best_train_epoch))
        print('Best test acc: {}, epoch: {}'.format(best_test_acc, best_test_epoch))
        print('best checkpoint:', best_checkpoint)

    elif args.test == True:
        if args.load is True:
            print('load data')
            test_data = torch.load('./data/DvsGesture/test_data_max.pt')
        else:
            print('process data')
            test_data = GestureDataset(root='./data/DvsGesture', train=False, length=length)
            torch.save(test_data, './data/DvsGesture/test_data_max.pt')
        print('test_data', len(test_data))
        test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False, drop_last=False)

        test_checkpoint = torch.load(test_checkpoint_path)
        snn.load_state_dict(test_checkpoint["snn_state_dict"])

        test_acc, test_loss = test(snn, test_dataloader)

        print('Test checkpoint: {}, acc: {}'.format(test_checkpoint_path, test_acc))