Deprecation Notice: This library was an experiment trying to get pytree Modules working with Flax-like colletions. I'd currently recommend the following alternatives:
A Pytree Module system for Deep Learning in JAX
jit
, grad
, vmap
, etc.Treex is implemented on top of Treeo and reexports all of its API for convenience.
Getting Started | User Guide | Examples | Documentation
Module
class.nn
module for with common layers implemented as wrappers over Flax layers.losses
module with common loss functions.metrics
module with common metrics.Optimizer
class that can wrap any optax optimizer.
Despite all JAX benefits, current Module systems are not intuitive to new users and add additional complexity not present in frameworks like PyTorch or Keras. Treex takes inspiration from S4TF and delivers an intuitive experience using JAX Pytree infrastructure.
Install using pip:
pip install treex
This is a small appetizer to give you a feel for how using Treex looks like, be sure to checkout the User Guide for a more in-depth explanation.
import treex as tx
import numpy as np
import jax, optax
# create some data
x = np.random.uniform(size=(50, 1))
y = 1.3 * x ** 2 - 0.3 + np.random.normal(size=x.shape)
# initialize a Module, its simple
model = tx.MLP([64, 1]).init(key=42, inputs=x)
# define an optimizer, init with model params
optimizer = tx.Optimizer(optax.adam(4e-3)).init(model)
# define loss function, notice
# Modules are jit-abel and differentiable π€―
@jax.jit
@jax.grad
def loss_fn(model: tx.MLP, x, y):
# forward is a simple call
preds = model(x)
# MSE
return ((preds - y) ** 2).mean()
# basic training loop
for step in range(500):
# grads have the same type as model
grads: tx.MLP = loss_fn(model, x, y)
# apply the gradient updates
model = optimizer.update(grads, model)
# Pytorch-like eval mode
model = model.eval()
preds = model(x)
Treex is in an early stage, things might break between versions but we will respect semanting versioning. Since Treex layers are numerically equivalent to Flax, it borrows some maturity and yields more confidence over its results. Feedback is much appreciated.
Roadmap:
Contributions are welcomed!
Checkout the /examples directory for more detailed examples. Here are a few additional toy examples:
This is a simple but realistic example of how Treex is used.
from functools import partial
from typing import Union
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import optax
import treex as tx
x = np.random.uniform(size=(500, 1))
y = 1.4 * x - 0.3 + np.random.normal(scale=0.1, size=(500, 1))
# differentiate only w.r.t. parameters
def loss_fn(params, model, x, y):
# merge params into model
model = model.merge(params)
preds = model(x)
loss = jnp.mean((preds - y) ** 2)
# the model may contain state updates
# so it should be returned
return loss, model
grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
# both model and optimizer are jit-able
@jax.jit
def train_step(model, x, y, optimizer):
# select only the parameters
params = model.parameters()
(loss, model), grads = grad_fn(params, model, x, y)
# update params and model
params = optimizer.update(grads, params)
model = model.merge(params)
# return new model and optimizer
return loss, model, optimizer
model = tx.Linear(1).init(42, x)
optimizer = tx.Optimizer(optax.adam(0.01)).init(model)
for step in range(300):
loss, model, optimizer = train_step(model, x, y, optimizer)
if step % 50 == 0:
print(f"loss: {loss:.4f}")
# eval mode "turns off" layers like Dropout / BatchNorm
model = model.eval()
X_test = np.linspace(x.min(), x.max(), 100)[:, None]
preds = model(X_test)
plt.scatter(x, y, c="k", label="data")
plt.plot(X_test, preds, c="b", linewidth=2, label="prediction")
plt.legend()
plt.show()
Here is an example of creating a stateful module of a RollingMean
metric and using them with jax.jit
. For a real use cases use the metrics inside treex.metrics
.
class RollingMean(tx.Module):
count: jnp.ndarray = tx.State.node()
total: jnp.ndarray = tx.State.node()
def __init__(self):
self.count = jnp.array(0, dtype=jnp.int32)
self.total = jnp.array(0.0, dtype=jnp.float32)
def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
self.count += np.prod(x.shape)
self.total += x.sum()
return self.total / self.count
@jax.jit
def update(x: jnp.ndarray, metric: RollingMean) -> Tuple[jnp.ndarray, RollingMean]:
mean = metric(x)
return mean, metric # return mean value and updated metric
metric = RollingMean()
for i in range(10):
x = np.random.uniform(-1, 1, size=(100, 1))
mean, metric = update(x, metric)
print(mean)