fastxtend

Train fastai models faster (and other useful tools)

fastxtend accelerates fastai

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 with CutMixUp or CutMixUpAugment.
• 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 into TensorSpec spectrograms or TensorMelSpec mel spectrograms using the GPU.
• Out of the box support for converting one TensorAudio to one or multiple TensorSpec or TensorMelSpec objects from the Datablock api.
• Audio MixUp and CutMix Callbacks.
• audio_learner which merges multiple TensorSpec or TensorMelSpec 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 pytorch -c nvidia/label/cuda-11.8.0 -c fastai \
-c huggingface -c conda-forge

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:

Learner(..., opt_func=adam(foreach=True))

Log an accuracy metric on the training set as a smoothed metric and validation set like normal:

Learner(..., metrics=[Accuracy(log_metric=LogMetric.Train, metric_type=MetricType.Smooth),
                      Accuracy()])

Log multiple losses as individual metrics on train and valid:

mloss = MultiLoss(loss_funcs=[nn.MSELoss, nn.L1Loss],
                  weights=[1, 3.5], loss_names=['mse_loss', 'l1_loss'])

Learner(..., loss_func=mloss, metrics=RMSE(), cbs=MultiLossCallback)

Apply MixUp, CutMix, or Augmentation while training:

Learner(..., cbs=CutMixUpAugment)

Profile a fastai training loop:

from fastxtend.callback import simpleprofiler

learn = Learner(...).profile()
learn.fit_one_cycle(2, 3e-3)

Benchmark

To run the benchmark on your own machine, see the example scripts for details on how to replicate.