Skip to content

Commit 1aa1145

Browse files
mattdangerwmattdangerw
committed
Update requirements and install instructions for multi-backend keras
1 parent 2c8e915 commit 1aa1145

8 files changed

+90
-83
lines changed

‎CONTRIBUTING.md‎

Lines changed: 36 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ Once the pull request is approved, a team member will take care of merging.
8181

8282
## Setting up an Environment
8383

84-
Python 3.7 or later is required.
84+
Python 3.9 or later is required.
8585

8686
Setting up your KerasNLP development environment requires you to fork the
8787
KerasNLP repository and clone it locally. With the
@@ -93,72 +93,53 @@ cd keras-nlp
9393
```
9494

9595
Next we must setup a python environment with the correct dependencies. We
96-
recommend using `conda` to install tensorflow dependencies (such as CUDA), and
97-
`pip` to install python packages from PyPI. The exact method will depend on your
98-
OS.
99-
100-
**Note**: Please be careful not to use the `tensorflow` pre-packaged with conda,
101-
which is incompatible with `tensorflow-text` on PyPi, and follow the
102-
instructions below.
96+
recommend using `conda` to set up a base environment, and `pip` to install
97+
python packages from PyPI. The exact method will depend on your OS.
10398

10499
### Linux (recommended)
105100

106-
To setup a complete environment with TensorFlow, a local install of keras-nlp,
107-
and all development tools, run the following or adapt it to suit your needs.
101+
For devloping and unit testing the library, a CPU only environment is often
102+
sufficient. For any training or inference with the library, you will quickly
103+
want accelerator support. The easiest way to get GPU support across all of our
104+
backends is to set up a few different python environements and pull in all cuda
105+
dependencies via `pip`.
106+
107+
The shell snippet below will install four conda environments: `keras-nlp-cpu`,
108+
`keras-nlp-jax`, `keras-nlp-torch`, and `keras-nlp-tensorflow`. The cpu
109+
environement supports all backends without cuda, and each backend environement
110+
has cuda support.
108111

109112
```shell
110-
# Create and activate conda environment.
111-
conda create -n keras-nlp python=3.9
112-
conda activate keras-nlp
113-
114-
# The following can be omitted if GPU support is not required.
115-
conda install -c conda-forge cudatoolkit-dev=11.2 cudnn=8.1.0
116-
mkdir -p $CONDA_PREFIX/etc/conda/activate.d/
117-
echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
118-
echo 'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
119-
source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
120-
121-
# Install dependencies.
122-
python -m pip install --upgrade pip
123-
python -m pip install -r requirements.txt
124-
python -m pip install -e "."
113+
conda create -y -n keras-nlp-cpu python=3.10
114+
conda activate keras-nlp-cpu
115+
pip install -r requirements.txt # install deps
116+
python pip_build.py --install # install keras-nlp
117+
118+
for backend in "jax" "torch" "tensorflow"; do
119+
conda create -y -n keras-nlp-${backend} python=3.10
120+
conda activate keras-nlp-${backend}
121+
pip install -r requirements-${backend}-cuda.txt # install deps
122+
python pip_build.py --install # install keras-nlp
123+
done
125124
```
126125

127-
### MacOS
128-
129-
⚠️⚠️⚠️ MacOS binaries are for the M1 architecture are not currently available from
130-
official sources. You can try experimental development workflow leveraging the
131-
[tensorflow metal plugin](https://developer.apple.com/metal/tensorflow-plugin/)
132-
and a [community maintained build](https://github.com/sun1638650145/Libraries-and-Extensions-for-TensorFlow-for-Apple-Silicon)
133-
of `tensorflow-text`. These binaries are not provided by Google, so proceed at
134-
your own risk.
135-
136-
#### Experimental instructions for Arm (M1)
126+
To activate the jax environment and set keras to use jax, run:
137127

138128
```shell
139-
# Create and activate conda environment.
140-
conda create -n keras-nlp python=3.9
141-
conda activate keras-nlp
142-
143-
# Install dependencies.
144-
conda install -c apple tensorflow-deps=2.9
145-
python -m pip install --upgrade pip
146-
python -m pip install -r requirements-macos-m1.txt
147-
python -m pip install -e "."
129+
conda activate keras-nlp-jax && export KERAS_BACKEND=jax
148130
```
149131

150-
#### Instructions for x86 (Intel)
132+
### MacOS
151133

152-
```shell
153-
# Create and activate conda environment.
154-
conda create -n keras-nlp python=3.9
155-
conda activate keras-nlp
156-
157-
# Install dependencies.
158-
python -m pip install --upgrade pip
159-
python -m pip install -r requirements.txt
160-
python -m pip install -e "."
161-
```
134+
`tensorflow-text` does not release precompiled binaries for MacOS M-series
135+
chips, though the library does support building from source on MacOS.
136+
137+
We strongly recommend a Linux development environment for an easy contribution
138+
experience. To build a dev environement from scratch on MacOS, see the following
139+
guides:
140+
141+
https://developer.apple.com/metal/tensorflow-plugin/
142+
https://github.com/tensorflow/text
162143

163144
### Windows
164145

‎requirements-common.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,3 +16,6 @@ namex
1616
# Optional deps.
1717
rouge-score
1818
sentencepiece
19+
tensorflow-datasets
20+
# Breakage fix.
21+
ml-dtypes==0.2.0

‎requirements-jax-cuda.txt‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
# Tensorflow (always required for preprocessing).
2+
tensorflow==2.14.0
3+
tensorflow-text==2.14.0.rc0
4+
5+
# Torch cpu-only version to assist with testing.
6+
--extra-index-url https://download.pytorch.org/whl/cpu
7+
torch==2.0.1
8+
torchvision==0.15.2
9+
10+
# Jax.
11+
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
12+
jax[cuda11_pip]==0.4.14
13+
14+
-r requirements-common.txt

‎requirements-macos-m1.txt‎

Lines changed: 0 additions & 16 deletions
This file was deleted.

‎requirements-nightly.txt‎

Lines changed: 0 additions & 7 deletions
This file was deleted.

‎requirements-tensorflow-cuda.txt‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Tensorflow with cuda support.
2+
tensorflow[and-cuda]==2.14.0
3+
tensorflow-text[and-cuda]==2.14.0.rc0
4+
5+
# Torch cpu-only version to assist with testing.
6+
--extra-index-url https://download.pytorch.org/whl/cpu
7+
torch==2.0.1
8+
torchvision==0.15.2
9+
10+
# Jax cpu-only version to assist with testing.
11+
jax[cpu]==0.4.14
12+
13+
-r requirements-common.txt

‎requirements-torch-cuda.txt‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Tensorflow (always required for preprocessing).
2+
tensorflow==2.14.0
3+
tensorflow-text==2.14.0.rc0
4+
5+
# Torch with cuda support.
6+
--extra-index-url https://download.pytorch.org/whl/cu117
7+
torch==2.0.1+cu117
8+
torchvision==0.15.2+cu117
9+
10+
# Jax cpu-only version to assist with testing.
11+
jax[cpu]==0.4.14
12+
13+
-r requirements-common.txt

‎requirements.txt‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,13 @@
1-
# Core deps.
2-
tensorflow~=2.13.0
3-
tensorflow-text~=2.13.0
4-
tensorflow-datasets
1+
# Tensorflow.
2+
tensorflow==2.14.0
3+
tensorflow-text==2.14.0.rc0
4+
5+
# Torch.
6+
--extra-index-url https://download.pytorch.org/whl/cpu
7+
torch==2.0.1
8+
torchvision==0.15.2
9+
10+
# Jax.
11+
jax[cpu]==0.4.14
512

6-
# Common deps.
713
-r requirements-common.txt

0 commit comments

Comments
 (0)