Elegy
Elegy is a framework-agnostic Trainer interface for the Jax ecosystem.
Main Features
- Easy-to-use: Elegy provides a Keras-like high-level API that makes it very easy to do common tasks.
- Flexible: Elegy provides a functional Pytorch Lightning-like low-level API that provides maximal flexibility when needed.
- Agnostic: Elegy supports a variety of frameworks including Flax, Haiku, and Optax on the high-level API, and it is 100% framework-agnostic on the low-level API.
- Compatible: Elegy can consume a wide variety of common data sources including TensorFlow Datasets, Pytorch DataLoaders, Python generators, and Numpy pytrees.
For more information take a look at the Documentation.
Installation
Install Elegy using pip:
pip install elegy
For Windows users we recommend the Windows subsystem for linux 2 WSL2 since jax does not support it yet.
Quick Start: High-level API
Elegy's high-level API provides a very simple interface you can use by implementing following steps:
1. Define the architecture inside a Module. We will use Flax Linen for this example:
import flax.linen as nn
import jax
class MLP(nn.Module):
@nn.compact
def call(self, x):
x = nn.Dense(300)(x)
x = jax.nn.relu(x)
x = nn.Dense(10)(x)
return x
2. Create a Model from this module and specify additional things like losses, metrics, and optimizers:
import elegy, optax
model = elegy.Model(
module=MLP(),
loss=[
elegy.losses.SparseCategoricalCrossentropy(from_logits=True),
elegy.regularizers.GlobalL2(l=1e-5),
],
metrics=elegy.metrics.SparseCategoricalAccuracy(),
optimizer=optax.rmsprop(1e-3),
)
3. Train the model using the fit method:
model.fit(
x=X_train,
y=y_train,
epochs=100,
steps_per_epoch=200,
batch_size=64,
validation_data=(X_test, y_test),
shuffle=True,
callbacks=[elegy.callbacks.TensorBoard("summaries")]
)
Quick Start: Low-level API
In Elegy's low-level API lets you define exactly what goes on during training, testing, and inference. Lets define the test_step to implement a linear classifier in pure jax:
1. Calculate our loss, logs, and states:
class LinearClassifier(elegy.Model):
# request parameters by name via depending injection.
# names: x, y_true, sample_weight, class_weight, states, initializing
def test_step(
self,
x, # inputs
y_true, # labels
states: elegy.States, # model state
initializing: bool, # if True we should initialize our parameters
):
rng: elegy.RNGSeq = states.rng
# flatten + scale
x = jnp.reshape(x, (x.shape[0], -1)) / 255
# initialize or use existing parameters
if initializing:
w = jax.random.uniform(
rng.next(), shape=[np.prod(x.shape[1:]), 10]
)
b = jax.random.uniform(rng.next(), shape=[1])
else:
w, b = states.net_params
# model
logits = jnp.dot(x, w) + b
# categorical crossentropy loss
labels = jax.nn.one_hot(y_true, 10)
loss = jnp.mean(-jnp.sum(labels * jax.nn.log_softmax(logits), axis=-1))
accuracy=jnp.mean(jnp.argmax(logits, axis=-1) == y_true)
# metrics
logs = dict(
accuracy=accuracy,
loss=loss,
)
return loss, logs, states.update(net_params=(w, b))
2. Instantiate our LinearClassifier with an optimizer:
model = LinearClassifier(
optimizer=optax.rmsprop(1e-3),
)
3. Train the model using the fit method:
model.fit(
x=X_train,
y=y_train,
epochs=100,
steps_per_epoch=200,
batch_size=64,
validation_data=(X_test, y_test),
shuffle=True,
callbacks=[elegy.callbacks.TensorBoard("summaries")]
)
Using Jax Frameworks
It is straightforward to integrate other functional JAX libraries with this low-level API:
class LinearClassifier(elegy.Model):
def test_step(
self, x, y_true, states: elegy.States, initializing: bool
):
rng: elegy.RNGSeq = states.rng
x = jnp.reshape(x, (x.shape[0], -1)) / 255
if initializing:
logits, variables = self.module.init_with_output(
{"params": rng.next(), "dropout": rng.next()}, x
)
else:
variables = dict(params=states.net_params, **states.net_states)
logits, variables = self.module.apply(
variables, x, rngs={"dropout": rng.next()}, mutable=True
)
net_states, net_params = variables.pop("params")
labels = jax.nn.one_hot(y_true, 10)
loss = jnp.mean(-jnp.sum(labels * jax.nn.log_softmax(logits), axis=-1))
accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == y_true)
logs = dict(accuracy=accuracy, loss=loss)
return loss, logs, states.update(net_params=net_params, net_states=net_states)
More Info
- Getting Started: High-level API tutorial.
- Getting Started: Low-level API tutorial.
- Elegy's Documentation.
- The examples directory.
- What is Jax?
Examples
To run the examples first install some required packages:
pip install -r examples/requirements.txt
Now run the example:
python examples/flax_mnist_vae.py
Contributing
Deep Learning is evolving at an incredible pace, there is so much to do and so few hands. If you wish to contribute anything from a loss or metric to a new awesome feature for Elegy just open an issue or send a PR! For more information check out our Contributing Guide.
About Us
We are some friends passionate about ML.
License
Apache
Citing Elegy
To cite this project:
BibTeX
@software{elegy2020repository,
author = {PoetsAI},
title = {Elegy: A framework-agnostic Trainer interface for the Jax ecosystem},
url = {https://github.com/poets-ai/elegy},
version = {0.5.0},
year = {2020},
}
Where the current version may be retrieved either from the Release tag or the file elegy/__init__.py and the year corresponds to the project's release year.
![[Bug] Problem with computing metrics](https://avatars.githubusercontent.com/u/29305032?v=4)
![Documentation/API reference not accessible via project website[Bug]](https://avatars.githubusercontent.com/u/62201809?v=4)
![[Bug] elegy does not work with latest haiku version](https://avatars.githubusercontent.com/u/4632336?v=4)
![[Bug] Elegy crash on GPU](https://avatars.githubusercontent.com/u/1382326?v=4)
53 Jan 01, 2023
48 Aug 19, 2022
21 Dec 14, 2022
140 Dec 19, 2022
25 Jan 09, 2023
1 Nov 11, 2022
8 Oct 17, 2022
5 Jun 02, 2022
0 Mar 20, 2022
22 Nov 09, 2022
3 Mar 26, 2022
88 Nov 22, 2022
62 Nov 21, 2022
35 Nov 25, 2022
255 Sep 07, 2022
430 Dec 23, 2022
63 Nov 22, 2022
4.7k Dec 29, 2022
46 Dec 26, 2022
142 Jan 01, 2023