AIToolbox - Model Training Framework for PyTorch

AIToolbox is a framework which helps you train deep learning models in PyTorch and quickly iterate experiments. It hides the repetitive technicalities of training the neural nets and frees you to focus on interesting part of devising new models. In essence, it offers a keras-style train loop abstraction which can be used for higher level training process while still allowing the manual control on the lower level when desired.

In addition to orchestrating the model training loop the framework also helps you keep track of different experiments by automatically saving models in a structured traceable way and creating performance reports. These can be stored both locally or on AWS S3 (Google Cloud Storage in beta) which makes the library very useful when training on the GPU instance on AWS. Instance can be automatically shut down when training is finished and all the results are safely stored on S3.

Main Components

AIToolbox consists of three main user-facing components:

All three AIToolbox components can be used independently when only some subset of functionality is desired in a project. However, the greatest benefit of AIToolbox comes when all components are used together in unison in order to ease the process of PyTorch model training and experiment tracking as much as possible. Most of this top-level API is exposed to the user via the functionality implemented in aitoolbox.torchtrain.


To install the AIToolbox package execute:

pip install aitoolbox

If you want to install the most recent version from github repository, first clone the package repository and then install via the pip command:

git clone

pip install ./aitoolbox

AIToolbox package can be also provided as a dependency in the requirements.txt file. This can be done by just specifying the aitoolbox dependency. On the other hand, to automatically download the current master branch from github include the following dependency specification in the requirements.txt:


Documentation Sections:

To learn more about components available in AIToolbox have a look at the corresponding documentation sections: