Objax is a machine learning framework that provides an Object Oriented layer for JAX.
Project description
Objax
Tutorials | Install | Documentation | Philosophy
This is not an officially supported Google product.
Objax is an open source machine learning framework that accelerates research and learning thanks to a minimalist object-oriented design and a readable code base. Its name comes from the contraction of Object and JAX -- a popular high-performance framework. Objax is designed by researchers for researchers with a focus on simplicity and understandability. Its users should be able to easily read, understand, extend, and modify it to fit their needs.
This is the developer repository of Objax, there is very little user documentation here, for the full documentation go to objax.readthedocs.io.
You can find READMEs in the subdirectory of this project, for example:
User installation guide
You install Objax using pip
as follows:
pip install --upgrade objax
Objax supports GPUs but assumes that you already have some version of CUDA installed. Here are the extra steps required to install CUDA-enabled jaxlib (jaxlib releases require CUDA 11.2 or newer):
RELEASE_URL="https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
JAX_VERSION=`python3 -c 'import jax; print(jax.__version__)'`
pip uninstall -y jaxlib
pip install -f $RELEASE_URL jax[cuda]==$JAX_VERSION
For more installation options, see https://github.com/google/jax#pip-installation-gpu-cuda
Useful environment configurations
Here are a few useful options:
# Prevent JAX from taking the whole GPU memory
# (useful if you want to run several programs on a single GPU)
export XLA_PYTHON_CLIENT_PREALLOCATE=false
Testing your installation
You can test your installation by running the code below:
import jax
import objax
print(f'Number of GPUs {jax.device_count()}')
x = objax.random.normal(shape=(100, 4))
m = objax.nn.Linear(nin=4, nout=5)
print('Matrix product shape', m(x).shape) # (100, 5)
x = objax.random.normal(shape=(100, 3, 32, 32))
m = objax.nn.Conv2D(nin=3, nout=4, k=3)
print('Conv2D return shape', m(x).shape) # (100, 4, 32, 32)
Typically if you get errors running this using CUDA, it probably means your installation of CUDA or CuDNN has issues.
Runing code examples
Clone the code repository:
git clone https://github.com/google/objax.git
cd objax/examples
Citing Objax
To cite this repository:
@software{objax2020github,
author = {{Objax Developers}},
title = {{Objax}},
url = {https://github.com/google/objax},
version = {1.2.0},
year = {2020},
}
Developer documentation
Here is information about development setup and a guide on adding new code.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
File details
Details for the file objax-1.8.0.tar.gz
.
File metadata
- Download URL: objax-1.8.0.tar.gz
- Upload date:
- Size: 59.9 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.11.5
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | 55d34f393ac9e57f9ff73be4d7c2d9efdb0ba8fb2c27d7564f2bd06fa19e90da |
|
MD5 | e845d1e4f1daa350483315248768a4f8 |
|
BLAKE2b-256 | 200ccbadc1a538f5ad3baa529b9acb75f5433d291a5bc23787ed239bb5b24d29 |
File details
Details for the file objax-1.8.0-py3-none-any.whl
.
File metadata
- Download URL: objax-1.8.0-py3-none-any.whl
- Upload date:
- Size: 86.7 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.11.5
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | 7b0e3b28751529adc1402a05cc1e71f6e01b1476352d44cbdc4094d2586d1d1c |
|
MD5 | c6fc0d1983cfe3b0fe41d9cada53bb18 |
|
BLAKE2b-256 | d4664847bec2b72c61d4d131e0dbff9ecf1b9eced084441e0cdb8a4f162738b3 |