umangjpatel/kerax

Keras-like APIs for JAX framework

Kerax

Keras-like APIs for the JAX library.

Features

• Enables high-performance machine learning research.
• Built-in support of popular optimization algorithms and activation functions.
• Runs seamlessly on CPU, GPU and even TPU! without any manual configuration required.

Quickstart

Code

from kerax.datasets import binary_tiny_mnist
from kerax.layers import Dense, Relu, Sigmoid
from kerax.losses import BCELoss
from kerax.metrics import binary_accuracy
from kerax.models import Sequential
from kerax.optimizers import SGD

data = binary_tiny_mnist.load_dataset(batch_size=200)
model = Sequential([Dense(100), Relu, Dense(1), Sigmoid])
model.compile(loss=BCELoss, optimizer=SGD(step_size=0.003), metrics=[binary_accuracy])
model.fit(data=data, epochs=10)
model.save(file_name="model")

interp = model.get_interpretation()
interp.plot_losses()
interp.plot_accuracy()

Output

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Epoch 10: 100%|██████████| 10/10 [00:02<00:00,  3.82it/s, train_loss : 0.192 :: valid_loss : 0.202 :: train_binary_accuracy : 1.000 :: valid_binary_accuracy : 1.000]

Process finished with exit code 0

Documentation (Coming soon...)

Developer's Notes

This project is developed and maintained by Umang Patel

kerax's Issues

Regularization layers

Description
Add regularization layers such as Inverted Dropout and Batch Norm

Forward prop code abstraction

Description
Refactor the forward prop code to accommodate different layers in the future.

Solution
Write forward prop code in layers API

Additional context
Refer to https://jax.readthedocs.io/en/latest/index.html

DataBunch-like API

Description
Implement a databunch-like API for loading image data.

Solution
Use PyTorch and TensorFlow data loading libraries

Additional context
Refer to jax.readthedocs.io for more info.

PyTorch implementation slow. Bring back JAX

Description
PyTorch support inefficient. Takes too much time.

Describe the solution you'd like
Reimplement the library in JAX

Additional context
Refer to https://github.com/google/jax

Regression loss functions missing

Feature request description
Regression algorithms missing.

Solution proposed
Implement loss functions like MSE, RMSE, etc.

Additional context
Refer to docs.fast.ai, PyTorch and Keras docs for implementation

Constant accuracy in CNN

Description
Constant accuracy and large losses while using CNN APIs

Expected behaviour
Accuracy should increase while loss should decrease

JAX imports not working

Description
JAX imports displays NotImplementedError

Expected behavior
To use JAX-numpy in accordance with the original NumPy package.

Proposed Solution
Replace JAX computation engine with PyTorch package.

CNN APIs

Description
Implement CNN layers and relevant APIs

Solution
Use LAX APIs

Additional context
Refer to jax.readthedocs.io for more info.

dnet.py filled with lots of dependencies

Bug description
APIs too congested in dnet.py script

Additional context
Needs API cleanup from ground up

Documentation and experiments

Description
Update README and add experiments.

Solution
Create experiments or examples package to work

Weight initialization techniques

Description
Implement various weight initialization schemes

Solution
Use jax.nn module for implementation notes.

Multi-class classification

Description
Softmax classification algorithm

Stax-like layers API

Description
Stax-like layers API refactoring

Solution
Write every layer with apply_fun and init_fun

JAX vmap implementation

Description
Implement JAX vmap structure for batching

Solution
Use JAX vmap for batched gradient descent implementation.

Additional context
Refer to https://jax.readthedocs.io/en/latest/index.html

Support for saving and loading models

Is your feature request related to a problem? Please describe.
After training large models, we shall be able to save and reload them for further experiments.

Describe the solution you'd like
Models shall be able to save (serialize) and load (deserialize) items as follows :

• Model parameters
• Layers in models

Additional context
Refer to flax serialization logic for more info. It involves serializing and deserializing JAX PyTrees to MessagePack format.

Optimizers

Description
Add optimization algorithms

Solution
Use JAX experimental module

Multinomial / Multiclass classification implementation

Description of feature request
Multiclass / Multinomial classification problems not supported.

Solution proposed
Implement softmax function

Additional context
Refer to docs.fast.ai, PyTorch and Keras documentation for implementation details