Download
Trax is an end-to-end library for deep learning that focuses on clear code and speed. It is actively used and maintained in the Google Brain team. Run a pre-trained Transformer, create a translator in a few lines of code. Features and resources, API docs, where to talk to us, how to open an issue and more. Walkthrough, how Trax works, how to make new models and train on your own data. Trax includes basic models (like ResNet, LSTM, Transformer) and RL algorithms (like REINFORCE, A2C, PPO). It is also actively used for research and includes new models like the Reformer and new RL algorithms like AWR. Trax has bindings to a large number of deep learning datasets, including Tensor2Tensor and TensorFlow datasets. You can use Trax either as a library from your own python scripts and notebooks or as a binary from the shell, which can be more convenient for training large models. It runs without any changes on CPUs, GPUs and TPUs.
Features
- The basic units flowing through Trax models are tensors
- In Trax we want numpy operations to run very fast
- We also want to automatically compute gradients of functions on tensors
- Gradients can be calculated using trax.fastmath.grad
- Layers are basic building blocks of Trax models
- Models in Trax are built from layers most often using the Serial and Branch combinators
Project Samples
