Installation

Here we will outline how to install the zea package and its dependencies. Besides installation through PyPI, we also provide Docker images for easy development (see the Docker section below). For instructions on contributing and development installation, see Contributing.

conda create -n zea python=3.12 # 3.10, 3.11, or 3.12
conda activate zea
pip install zea
# ! note, you still need to install a backend
# see below for details

docker pull zeahub/all:latest
docker run --gpus 'all' -it zeahub/all:latest

Backend

zea is written in Python on top of Keras 3. This means that under the hood we use the Keras framework to implement the pipeline and models. Keras allows you to set a backend, which means you can use zea alongside your project that uses any of your preferred machine learning framework.

To use zea, you need to install one of the supported machine learning backends: JAX, PyTorch or TensorFlow zea will not run without a backend installed.

If you are unsure which backend to use, we recommend JAX as it is currently the fastest backend.

After installing a backend, set the KERAS_BACKEND environment variable to one of the following:

# at the top of your script before other imports
import os
os.environ["KERAS_BACKEND"] = "jax"
import zea

conda env config vars set KERAS_BACKEND=jax

export KERAS_BACKEND=jax

# at the top of your script before other imports
import os
os.environ["KERAS_BACKEND"] = "torch"
import zea

conda env config vars set KERAS_BACKEND=torch

export KERAS_BACKEND=torch

# at the top of your script before other imports
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import zea

conda env config vars set KERAS_BACKEND=tensorflow

export KERAS_BACKEND=tensorflow

# at the top of your script before other imports
# note NumPy backend has limited functionality
import os
os.environ["KERAS_BACKEND"] = "numpy"
import zea

# note NumPy backend has limited functionality
conda env config vars set KERAS_BACKEND=numpy

# note NumPy backend has limited functionality
export KERAS_BACKEND=numpy

Docker

This repository provides multiple Docker images built and tested in our CI pipeline.

Pre-built images

The following images are available on Docker Hub:

  • zeahub/all: This image includes support for all machine learning backends (TensorFlow, PyTorch, and JAX).

  • zeahub/tensorflow: This image includes support for TensorFlow.

  • zeahub/torch: This image includes support for PyTorch.

  • zeahub/jax: This image includes support for JAX.

These images are uploaded to Docker Hub via the CI pipeline and can be used directly in your projects via:

docker pull zeahub/all:latest

Build

One can build an image for a specific backend using the provided Dockerfile and build arguments.

docker build --build-arg INSTALL_JAX=gpu . -t zeahub/jax:latest

To build the full image including all backends with GPU support:

docker build --build-arg INSTALL_JAX=gpu --build-arg INSTALL_TORCH=gpu --build-arg INSTALL_TF=gpu . -t zeahub/all:latest

Note that these build arguments can be set to either cpu or gpu depending on your needs.

Run

Run a container with one of the built images. Ensure you mount your repository at /zea so that changes are reflected inside the container, and use your user and group IDs to avoid permission issues.

docker run --name {CONTAINER-NAME} --gpus 'all' \
  -v ~/zea:/zea \
  -d -it --user "$(id -u):$(id -g)" \
  {IMAGE-NAME}:{IMAGE-TAG}
Docker run command flags explained

  • docker run: create and run a new container from an image.

  • --name: name the container.

  • --gpus: specify GPU devices to add to the container (‘all’ to pass all GPUs).

  • -v or --volume: bind mount a volume.

  • -d or --detach: start the container as a background process.

  • -it: start an interactive terminal session.

    • --interactive: keep STDIN open.

    • --tty: allocate a pseudo-TTY.

  • -m or --memory: set a memory limit (use g for gigabytes).

  • --cpus: specify the number of CPU cores to use.

  • --user: run as a specific user.

  • -w or --workdir: set the working directory inside the container.

  • --rm: automatically remove the container when it exits.

  • --env-file: load environment variables from a .env file.

  • --hostname: set the container hostname (useful for users.yaml file).

Important

Mount your zea repository to /zea inside the container so that changes are reflected in the zea installation inside the container. Additionally, use your user ID and group ID with --user "$(id -u):$(id -g)" to avoid permission issues when writing to mounted volumes.

Tip

The Docker container sets a random hostname by default. You can set a hostname with the --hostname flag. This is useful for the users.yaml file. Alternatively, you can use the hostname wildcard in the users.yaml file.