pytorch-accelerated
is a lightweight library designed to accelerate the process of training PyTorch models by providing a minimal, but extensible training loop - encapsulated in a single Trainer
object - which is flexible enough to handle the majority of use cases, and capable of utilizing different hardware options with no code changes required.
pytorch-accelerated
offers a streamlined feature set, and places a huge emphasis on simplicity and transparency, to enable users to understand exactly what is going on under the hood, but without having to write and maintain the boilerplate themselves!
The key features are:
Significant effort has been taken to ensure that every part of the library - both internal and external components - is as clear and simple as possible, making it easy to customise, debug and understand exactly what is going on behind the scenes at each step; most of the behaviour of the trainer is contained in a single class! In the spirit of Python, nothing is hidden and everything is accessible.
pytorch-accelerated
is proudly and transparently built on top of Hugging Face Accelerate, which is responsible for the movement of data between devices and launching of training configurations. When customizing the trainer, or launching training, users are encouraged to consult the Accelerate documentation to understand all available options; Accelerate provides convenient functions for operations such gathering tensors and gradient clipping, usage of which can be seen in the pytorch-accelerated
examples folder!
To learn more about the motivations behind this library, along with a detailed getting started guide, check out this blog post.
pytorch-accelerated
can be installed from pip using the following command:
pip install pytorch-accelerated
To make the package as slim as possible, the packages required to run the examples are not included by default. To include these packages, you can use the following command:
pip install pytorch-accelerated[examples]
To get started, simply import and use the pytorch-accelerated Trainer
,as demonstrated in the following snippet, and then launch training using the accelerate CLI described below.
# examples/core/train_mnist.py import os from torch import nn, optim from torch.utils.data import random_split from torchvision import transforms from torchvision.datasets import MNIST from pytorch_accelerated import Trainer class MNISTModel(nn.Module): def __init__(self): super().__init__() self.main = nn.Sequential( nn.Linear(in_features=784, out_features=128), nn.ReLU(), nn.Linear(in_features=128, out_features=64), nn.ReLU(), nn.Linear(in_features=64, out_features=10), ) def forward(self, input): return self.main(input.view(input.shape[0], -1)) def main(): dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor()) train_dataset, validation_dataset, test_dataset = random_split(dataset, [50000, 5000, 5000]) model = MNISTModel() optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) loss_func = nn.CrossEntropyLoss() trainer = Trainer( model, loss_func=loss_func, optimizer=optimizer, ) trainer.train( train_dataset=train_dataset, eval_dataset=validation_dataset, num_epochs=8, per_device_batch_size=32, ) trainer.evaluate( dataset=test_dataset, per_device_batch_size=64, ) if __name__ == "__main__": main()
To launch training using the accelerate CLI , on your machine(s), run:
accelerate config --config_file accelerate_config.yaml
and answer the questions asked. This will generate a config file that will be used to properly set the default options when doing
accelerate launch --config_file accelerate_config.yaml train.py [--training-args]
Note: Using the accelerate CLI is completely optional, training can also be launched in the usual way using:
python train.py
/ python -m torch.distributed ...
depending on your infrastructure configuration, for users who would like to maintain a more fine-grained control over the launch command.
More complex training examples can be seen in the examples folder here.
Alternatively, if you would rather understand the core concepts first, this can be found in the documentation.
Who is pytorch-accelerated aimed at?pytorch-accelerated
focuses only on the training process, with all other concerns being left to the responsibility of the user.Trainer
, the training loop is fundamentally broken up into a number of different methods that you would have to override. But, before you go, is writing those for
loops really important enough to warrant starting from scratch again 😉.Many aspects behind the design and features of pytorch-accelerated
were greatly inspired by a number of excellent libraries and frameworks such as fastai, timm, PyTorch-lightning and Hugging Face Accelerate. Each of these tools have made an enormous impact on both this library and the machine learning community, and their influence can not be stated enough!
pytorch-accelerated
has taken only inspiration from these tools, and all of the functionality contained has been implemented from scratch in a way that benefits this library. The only exceptions to this are some of the scripts in the examples folder in which existing resources were taken and modified in order to showcase the features of pytorch-accelerated
; these cases are clearly marked, with acknowledgement being given to the original authors.
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4