add batch trainer: the trainer can directly read inputs from torch dataloader without to load .npy files. It saves a lot memory.
https://arxiv.org/abs/2105.00324
git clone https://github.com/byin-cwi/neko.git
cd neko
pip install -e .
Train a RSNN with ALIF neurons with e-prop on MNIST:
from neko.backend import pytorch_backend as backend
# from neko.datasets import MNIST
from neko.evaluator import Evaluator
from neko.layers import ALIFRNNModel
from neko.learning_rules import Eprop
from neko.trainers import Trainer
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torch
trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
# if not exist, download mnist dataset
train_set = dset.MNIST(root='./', train=True, transform=trans, download=True)
test_set = dset.MNIST(root='./', train=False, transform=trans, download=True)
batch_size = 100
train_loader = torch.utils.data.DataLoader(
dataset=train_set,
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(
dataset=test_set,
batch_size=batch_size,
shuffle=False)
print('==>>> total trainning batch number: {}'.format(len(train_loader)))
print('==>>> total testing batch number: {}'.format(len(test_loader)))
model = ALIFRNNModel(128, 10, backend=backend, task_type='classification', return_sequence=False)
evaluated_model = Evaluator(model=model, loss='categorical_crossentropy', metrics=['accuracy', 'firing_rate'])
algo = Eprop(evaluated_model, mode='symmetric')
trainer = Trainer(algo)
trainer.train(train_loader,input_size=[28,28], T = 28, n_classes=10, epochs=30)
Training on the MNIST dataset with the same setting above, but more options available. For example, you can learn with BPTT and three variations of e-prop.
python examples/mnist.py
Training on the TIMIT dataset. You need to place the timit_processed
folder the same place as the script containing the processed dataset
produced by a script from the original authors of e-prop.
python examples/timit.py
Regularization enabled:
python timit.py --reg --eprop_mode symmetric --reg_coeff 5e-7
# Test: {'loss': 0.8918977379798889, 'accuracy': 0.7501428091397849, 'firing_rate': 12.973159790039062}
Faster training (~7.5X, 28s per epoch with RTX3090) with regularization enabled:
python timit.py --reg --eprop_mode symmetric --reg_coeff 3e-8 --batch_size 256 --learning_rate 0.01
# Test: {'loss': 0.8605409860610962, 'accuracy': 0.7542506720430108, 'firing_rate': 13.105131149291992}
Training on the MNIST-1D dataset with HMC:
python examples/mnist_1d_hmc.py
Training on the MNIST dataset with the simple Manhattan rule or Mahattan material rule:
python examples/mnist_manhattan.py
Compare the gradients from BPTT with the three varients of e-prop:
python examples/mnist_gradcompare.py
This is a visualization from the results of the script above.