Trax — your path to advanced deep learning

Trax helps you understand deep learning. We start with basic maths and go through layers, models, supervised and reinforcement learning. We get to advanced deep learning results, including recent papers and state-of-the-art models.

Trax is a successor to the Tensor2Tensor library and is actively used and maintained by researchers and engineers within the Google Brain team and a community of users. We're eager to collaborate with you too, so feel free to open an issue on GitHub or send along a pull request (see our contribution doc).

Examples

See our example layers in a TPU/GPU-backed colab notebook at Trax Demo

MLP on MNIST

python -m trax.trainer \
  --dataset=mnist \
  --model=MLP \
  --config="train.train_steps=1000"

Resnet50 on Imagenet

python -m trax.trainer \
  --config_file=$PWD/trax/configs/resnet50_imagenet_8gb.gin

TransformerDecoder on LM1B

python -m trax.trainer \
  --config_file=transformer_lm1b_8gb.gin