A lightweight framework for Hebbian learning based on PyTorch Ignite. Accepted at the Beyond Backpropagation NeurIPS worskhop.
A lightweight and flexible framework for Hebbian learning in PyTorch.
To get a local copy up and running follow these simple steps.
Depending on your system (OS/GPU/CUDA support) you may need to manually install a specific PyTorch version. Please see the PyTorch website for more information.
Clone the repository:
pip install git+https://github.com/Joxis/pytorch-hebbian.git
Similar to a classic PyTorch pipeline, a model, dataset and data loader are first initiated. Next, the framework requires you to specify a learning rule, optimizer and trainer. You can then start Hebbian learning.
# Creating the model
model = models.create_fc1_model([28 ** 2, 400])
# Creating the dataset and data loaders
dataset = datasets.mnist.MNIST(root=config.DATASETS_DIR, download=True,
transform=transforms.ToTensor())
train_loader = DataLoader(dataset, batch_size=1024, shuffle=True)
# Creating the learning rule, optimizer and trainer
learning_rule = KrotovsRule()
optimizer = Local(named_params=model.named_parameters(), lr=0.01)
trainer = HebbianTrainer(model=model, learning_rule=learning_rule,
optimizer=optimizer)
# Running the trainer
trainer.run(train_loader=train_loader, epochs=10)
Various TensorBoard visualizations are supported. Below are some examples.
See examples for more.
Any contributions you make are greatly appreciated. For major changes, please open an issue first to discuss what you would like to change.
- Fork the project
- Create your feature branch (
git checkout -b feature/amazing-feature
) - Commit your changes (
git commit -m 'Add some amazing feature'
) - Push to the branch (
git push origin feature/amazing-feature
) - Open a pull request
Distributed under the MIT License. See LICENSE for more information.
Please consider citing this work if it helps your research.
@article{TODO}
Jules Talloen - <firstname>@<lastname>.eu
Project Link: https://github.com/Joxis/pytorch-hebbian