A High Level API for Deep Learning in JAX
Elegy is built on top of Treex and Treeo and reexports their APIs for convenience.
Getting Started | Examples | Documentation
Model
class with an Estimator-like API.callbacks
module with common Keras callbacks.From Treex
Module
class.nn
module for with common layers.losses
module with common loss functions.metrics
module with common metrics.Install using pip:
pip install elegy
For Windows users, we recommend the Windows subsystem for Linux 2 WSL2 since jax does not support it yet.
Elegy's high-level API provides a straightforward interface you can use by implementing the following steps:
1. Define the architecture inside a Module
:
import jax
import elegy as eg
class MLP(eg.Module):
@eg.compact
def __call__(self, x):
x = eg.Linear(300)(x)
x = jax.nn.relu(x)
x = eg.Linear(10)(x)
return x
2. Create a Model
from this module and specify additional things like losses, metrics, and optimizers:
import optax optax
import elegy as eg
model = eg.Model(
module=MLP(),
loss=[
eg.losses.Crossentropy(),
eg.regularizers.L2(l=1e-5),
],
metrics=eg.metrics.Accuracy(),
optimizer=optax.rmsprop(1e-3),
)
3. Train the model using the fit
method:
model.fit(
inputs=X_train,
labels=y_train,
epochs=100,
steps_per_epoch=200,
batch_size=64,
validation_data=(X_test, y_test),
shuffle=True,
callbacks=[eg.callbacks.TensorBoard("summaries")]
)
Elegy's low-level API lets you explicitly define what goes on during training, testing, and inference. Let's define our own custom Model
to implement a LinearClassifier
with pure JAX:
1. Define a custom init_step
method:
class LinearClassifier(eg.Model):
# use treex's API to declare parameter nodes
w: jnp.ndarray = eg.Parameter.node()
b: jnp.ndarray = eg.Parameter.node()
def init_step(self, key: jnp.ndarray, inputs: jnp.ndarray):
self.w = jax.random.uniform(
key=key,
shape=[features_in, 10],
)
self.b = jnp.zeros([10])
self.optimizer = self.optimizer.init(self)
return self
Here we declared the parameters w
and b
using Treex's Parameter.node()
for pedagogical reasons, however normally you don't have to do this since you typically use a sub-Module
instead.
2. Define a custom test_step
method:
def test_step(self, inputs, labels):
# flatten + scale
inputs = jnp.reshape(inputs, (inputs.shape[0], -1)) / 255
# forward
logits = jnp.dot(inputs, self.w) + self.b
# crossentropy loss
target = jax.nn.one_hot(labels["target"], 10)
loss = optax.softmax_cross_entropy(logits, target).mean()
# metrics
logs = dict(
acc=jnp.mean(jnp.argmax(logits, axis=-1) == labels["target"]),
loss=loss,
)
return loss, logs, self
3. Instantiate our LinearClassifier
with an optimizer:
model = LinearClassifier(
optimizer=optax.rmsprop(1e-3),
)
4. Train the model using the fit
method:
model.fit(
inputs=X_train,
labels=y_train,
epochs=100,
steps_per_epoch=200,
batch_size=64,
validation_data=(X_test, y_test),
shuffle=True,
callbacks=[eg.callbacks.TensorBoard("summaries")]
)
Check out the /example directory for some inspiration. To run an example, first install some requirements:
pip install -r examples/requirements.txt
And the run it normally with python e.g.
python examples/flax/mnist_vae.py
If your are interested in helping improve Elegy check out the Contributing Guide.
BibTeX
@software{elegy2020repository,
title = {Elegy: A High Level API for Deep Learning in JAX},
author = {PoetsAI},
year = 2021,
url = {https://github.com/poets-ai/elegy},
version = {0.8.1}
}