fastxtend
Train fastai models faster with fastxtend’s fused optimizers, Progressive Resizing callback, and integrated FFCV DataLoader.
Feature overview
Train Models Faster
- Drop in fused optimizers, which are 21 to 293 percent faster then fastai native optimizers.
- Up to 75% optimizer memory savings with integrated bitsandbytes 8-bit optimizers.
- Increase GPU throughput and decrease training time with the Progressive Resizing callback.
- Use the highly optimized FFCV DataLoader, fully integrated with fastai.
- Integrated support for
torch.compile
via the Compile callbacks.
General Features
- Fused implementations of modern optimizers, such as Adan and Lion.
- Flexible metrics which can log on train, valid, or both. Backwards compatible with fastai metrics.
- Easily use multiple losses and log each individual loss on train and valid.
- Multiple profilers for profiling training and identifying bottlenecks.
- A fast Exponential Moving Average callback for smoother training.
Vision
- Apply
MixUp
,CutMix
, or Augmentations at once withCutMixUp
orCutMixUpAugment
. - Additional image augmentations.
- Support for running fastai batch transforms on CPU.
- More attention and pooling modules
- A flexible implementation of fastai’s
XResNet
.
Audio
TensorAudio
,TensorSpec
,TensorMelSpec
objects which maintain metadata and support plotting themselves using librosa.- A selection of performant audio augmentations inspired by fastaudio and torch-audiomentations.
- Uses TorchAudio to quickly convert
TensorAudio
waveforms intoTensorSpec
spectrograms orTensorMelSpec
mel spectrograms using the GPU. - Out of the box support for converting one
TensorAudio
to one or multipleTensorSpec
orTensorMelSpec
objects from the Datablock api. - Audio MixUp and CutMix Callbacks.
audio_learner
which merges multipleTensorSpec
orTensorMelSpec
objects before passing to the model.
Check out the documentation for additional splitters, callbacks, schedulers, utilities, and more.
Install
fastxtend is avalible on pypi:
pip install fastxtend
To install with dependencies for vision, FFCV, audio, or all tasks run one of:
pip install fastxtend[vision]
pip install fastxtend[ffcv]
pip install fastxtend[audio]
pip install fastxtend[all]
Or to create an editable development install:
git clone https://github.com/warner-benjamin/fastxtend.git
cd fastxtend
pip install -e ".[dev]"
To easily install prerequisites for all fastxtend features, use Conda or Miniconda:
conda create -n fastxtend python=3.10 "pytorch>=2.0.0" \
\
torchvision torchaudio pytorch-cuda=11.8 cuda fastai nbdev \
pkg-config libjpeg-turbo opencv tqdm terminaltables psutil \
numpy numba librosa=0.9.2 timm kornia rich typer wandb -c nvidia/label/cuda-11.8.0 -c fastai \
-c pytorch -c conda-forge
-c huggingface
conda activate fastxtend
replacing pytorch-cuda=11.8
and nvidia/label/cuda-11.8.0
with your preferred supported version of Cuda. Then install fastxtend using pip
:
pip install fastxtend[all]
Usage
Like fastai, fastxtend provides safe wildcard imports using python’s __all__
.
from fastai.vision.all import *
from fastxtend.vision.all import *
from fastxtend.ffcv.all import *
In general, import fastxtend after all fastai imports, as fastxtend modifies fastai. Any method modified by fastxtend is backwards compatible with the original fastai code.
Examples
Use a fused ForEach optimizer:
=adam(foreach=True)) Learner(..., opt_func
Log an accuracy metric on the training set as a smoothed metric and validation set like normal:
=[Accuracy(log_metric=LogMetric.Train, metric_type=MetricType.Smooth),
Learner(..., metrics Accuracy()])
Log multiple losses as individual metrics on train and valid:
= MultiLoss(loss_funcs=[nn.MSELoss, nn.L1Loss],
mloss =[1, 3.5], loss_names=['mse_loss', 'l1_loss'])
weights
=mloss, metrics=RMSE(), cbs=MultiLossCallback) Learner(..., loss_func
Apply MixUp, CutMix, or Augmentation while training:
=CutMixUpAugment) Learner(..., cbs
Profile a fastai training loop:
from fastxtend.callback import simpleprofiler
= Learner(...).profile()
learn 2, 3e-3) learn.fit_one_cycle(
Benchmark
To run the benchmark on your own machine, see the example scripts for details on how to replicate.