QDucasse / nn_benchmark

🧠 Benchmark facility to train networks on different datasets for PyTorch/Brevitas
MIT License
24 stars 1 forks source link

Neural Network benchmark

This project consists of an utility wrapper around PyTorch and Brevitas to specify the network, dataset and parameters to train with.

Usage

You can run the training of LeNet on the MNIST dataset as follows:

$ PYTORCH_JIT=1 python nn_benchmark/main.py --network LeNet --dataset MNIST --epochs 3

Or its quantized version QuantLeNet on the CIFAR-10 dataset with bit-width of (4,4,8) (corresponding to activation, weight and input) :

$ PYTORCH_JIT=1 python nn_benchmark/main.py --network QuantLeNet --dataset CIFAR10 --epochs 3 \
$ --acq 4 --weq 4 --inq 8

The results can be observed under the experiments folder.

You can then evaluate your network with the following command:

$ python nn_benchmark/main.py --network LeNet --dataset MNIST --evaluate --resume ./experiments/<your_folder>/checkpoints/best.tar

Available networks and datasets

The following networks are supported:

Their quantized counterparts are available as well:

The following datasets can be used:

Installation

I worked on the project through a virtual environment with virtualenvwrapper and I highly recommend to do so as well. However, whether or not you are in a virtual environment, the installation proceeds as follows:

Finally, whether you chose the first or second option, you will need brevitas if you want to use quantized networks. The installation is better performed from source and can be done as follows (in your native or virtual environment):

    $ git clone https://github.com/Xilinx/brevitas.git
    $ cd brevitas
    $ pip install -e .

Structure of the project

Quick presentation of the different modules of the project:

This project uses the following external libraries:

If installed as specified above, the requirements are stated in the requirements.txt file and therefore automatically installed.
However, you can install each of them separately with the command (except for brevitas, please follow the installation from source provided at the end of the installation paragraph):

  $ pip install <library>

Objectives and Milestones of the project

Testing

All tests are written to work with nose and/or pytest. Just type pytest or nosetests as a command line in the project. Every test file can still be launched by executing the testfile itself.

  $ python nn_benchmark/tests/chosentest.py
  $ pytest

References