From 5a1549d5cce8b9d4ccf52f265cdfe429c92ae1ae Mon Sep 17 00:00:00 2001 From: Huguens Jean <40362353+huguensjean@users.noreply.github.com> Date: Mon, 2 Oct 2023 12:11:16 -0400 Subject: [PATCH] Open source jax vision transformer to model garden. (#2295) --- .../jax_vit_model_conversion.Dockerfile | 151 ++++++++++++++++++ .../dockerfile/train_vit_gpu.Dockerfile | 149 +++++++++++++++++ .../vit_config_without_data.py | 27 ++++ .../jax_vision_transformer/vit_jax2tf.py | 0 4 files changed, 327 insertions(+) create mode 100644 community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/dockerfile/jax_vit_model_conversion.Dockerfile create mode 100644 community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/dockerfile/train_vit_gpu.Dockerfile create mode 100644 community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/vit_config_without_data.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/vit_jax2tf.py diff --git a/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/dockerfile/jax_vit_model_conversion.Dockerfile b/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/dockerfile/jax_vit_model_conversion.Dockerfile new file mode 100644 index 000000000..9b8b868f8 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/dockerfile/jax_vit_model_conversion.Dockerfile @@ -0,0 +1,151 @@ +# This Dockerfile converts JAX vision transformer model to +# tensorflow saved model format. +# Here is an example to build this dockerfile: +# PROJECT="your gcp project" +# IMAGE_TAG="jax-vit-model-conversion:${USER}-test" +# docker build -f model_oss/jax_vision_transformer/dockerfile/jax_vit_model_conversion.Dockerfile . -t "${IMAGE_TAG}" +# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}" +# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}" + + +FROM tensorflow/tensorflow:2.12.0-gpu + +ENV DEBIAN_FRONTEND=noninteractive + +# Install basic libs +RUN apt-get update && apt-get install -y --no-install-recommends \ + cmake \ + curl \ + wget \ + sudo \ + gnupg \ + libsm6 \ + libxext6 \ + libxrender-dev \ + lsb-release \ + ca-certificates \ + build-essential \ + git + +# Copy Apache license. +RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE + +# Get 'vision_transformer' repository from github. +RUN git clone https://github.com/google-research/vision_transformer +# Set current directory to the downloaded 'vision_transformer' repository. +WORKDIR ./vision_transformer +# Using git reset command to pin it down to a specific version. +RUN git reset --hard e66b4732d44504251197a3da3f5949f3f3ce9ca6 + +# Install required libs +RUN pip install --upgrade pip +# The following pip installs are pinned down versions of those inside +# vit_jax/requirements.txt file. +# NOTE: Using `no-deps` flag to avoid overwriting of +# dependent library versions. For example, +# both `chex` and `jax` can overwrite each others +# `jax-lib` version. +RUN pip install --no-deps absl-py==1.4.0 +RUN pip install --no-deps aqtp==0.0.10 +RUN pip install --no-deps array-record==0.2.0 +RUN pip install --no-deps astunparse==1.6.3 +RUN pip install --no-deps cached-property==1.5.2 +RUN pip install --no-deps cachetools==5.3.0 +RUN pip install --no-deps certifi==2019.11.28 +RUN pip install --no-deps chardet==3.0.4 +RUN pip install --no-deps chex==0.1.7 +RUN pip install --no-deps click==8.1.3 +RUN pip install --no-deps cloudpickle==2.2.1 +RUN pip install --no-deps clu==0.0.9 +RUN pip install --no-deps contextlib2==21.6.0 +RUN pip install --no-deps dacite==1.8.1 +RUN pip install --no-deps dbus-python==1.2.16 +RUN pip install --no-deps decorator==5.1.1 +RUN pip install --no-deps dm-tree==0.1.8 +RUN pip install --no-deps einops==0.6.1 +RUN pip install --no-deps etils==1.3.0 +RUN pip install --no-deps flatbuffers==23.3.3 +RUN pip install --no-deps flax==0.6.10 +RUN pip install --no-deps git+https://github.com/google/flaxformer@9adaa4467cf17703949b9f537c3566b99de1b416 +RUN pip install --no-deps gast==0.4.0 +RUN pip install --no-deps google-auth==2.16.2 +RUN pip install --no-deps google-auth-oauthlib==0.4.6 +RUN pip install --no-deps google-pasta==0.2.0 +RUN pip install --no-deps googleapis-common-protos==1.59.0 +RUN pip install --no-deps grpcio==1.51.3 +RUN pip install --no-deps h5py==3.8.0 +RUN pip install --no-deps idna==2.8 +RUN pip install --no-deps importlib-metadata==6.1.0 +RUN pip install --no-deps importlib-resources==5.12.0 +RUN pip install --no-deps keras==2.12.0 +RUN pip install --no-deps libclang==16.0.0 +RUN pip install --no-deps Markdown==3.4.3 +RUN pip install --no-deps markdown-it-py==2.2.0 +RUN pip install --no-deps MarkupSafe==2.1.2 +RUN pip install --no-deps mdurl==0.1.2 +RUN pip install --no-deps ml-collections==0.1.1 +RUN pip install --no-deps msgpack==1.0.5 +RUN pip install --no-deps nest-asyncio==1.5.6 +RUN pip install --no-deps numpy==1.23.5 +RUN pip install --no-deps oauthlib==3.2.2 +RUN pip install --no-deps opt-einsum==3.3.0 +RUN pip install --no-deps optax==0.1.5 +RUN pip install --no-deps orbax-checkpoint==0.1.6 +RUN pip install --no-deps packaging==23.0 +RUN pip install --no-deps pandas==2.0.1 +RUN pip install --no-deps pip==23.1.2 +RUN pip install --no-deps promise==2.3 +RUN pip install --no-deps protobuf==4.22.1 +RUN pip install --no-deps psutil==5.9.5 +RUN pip install --no-deps pyasn1==0.4.8 +RUN pip install --no-deps pyasn1-modules==0.2.8 +RUN pip install --no-deps Pygments==2.15.1 +RUN pip install --no-deps PyGObject==3.36.0 +RUN pip install --no-deps python-apt==2.0.1+ubuntu0.20.4.1 +RUN pip install --no-deps python-dateutil==2.8.2 +RUN pip install --no-deps pytz==2023.3 +RUN pip install --no-deps PyYAML==6.0 +RUN pip install --no-deps requests==2.22.0 +RUN pip install --no-deps requests-oauthlib==1.3.1 +RUN pip install --no-deps requests-unixsocket==0.2.0 +RUN pip install --no-deps rich==13.3.5 +RUN pip install --no-deps rsa==4.9 +RUN pip install --no-deps scipy==1.10.1 +RUN pip install --no-deps setuptools==67.6.0 +RUN pip install --no-deps six==1.14.0 +RUN pip install --no-deps tensorboard==2.12.0 +RUN pip install --no-deps tensorboard-data-server==0.7.0 +RUN pip install --no-deps tensorboard-plugin-wit==1.8.1 +RUN pip install --no-deps tensorflow==2.12.0 +RUN pip install --no-deps tensorflow-cpu==2.12.0 +RUN pip install --no-deps tensorflow-datasets==4.9.2 +RUN pip install --no-deps tensorflow-estimator==2.12.0 +RUN pip install --no-deps tensorflow-hub==0.13.0 +RUN pip install --no-deps tensorflow-io-gcs-filesystem==0.31.0 +RUN pip install --no-deps tensorflow-metadata==1.13.1 +RUN pip install --no-deps tensorflow-probability==0.20.0 +RUN pip install --no-deps tensorflow-text==2.12.1 +RUN pip install --no-deps tensorstore==0.1.36 +RUN pip install --no-deps termcolor==2.2.0 +RUN pip install --no-deps toml==0.10.2 +RUN pip install --no-deps toolz==0.12.0 +RUN pip install --no-deps tqdm==4.65.0 +RUN pip install --no-deps typing_extensions==4.5.0 +RUN pip install --no-deps tzdata==2023.3 +RUN pip install --no-deps urllib3==1.25.8 +RUN pip install --no-deps Werkzeug==2.2.3 +RUN pip install --no-deps wheel==0.40.0 +RUN pip install --no-deps wrapt==1.14.1 +RUN pip install --no-deps zipp==3.15.0 +# Installing jax at the very end with GPU support. +# NOTE: Not using `no-deps` flag here because +# we need CUDA support. +RUN pip install jax[cuda11_cudnn82]==0.4.6 \ +--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html + +ENV PYTHONPATH ./vit_jax + +COPY ./model_oss/jax_vision_transformer/vit_jax2tf.py ./ +COPY ./model_oss/jax_vision_transformer/vit_config_without_data.py vit_jax/configs/vit.py + +ENTRYPOINT ["python", "vit_jax2tf.py"] \ No newline at end of file diff --git a/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/dockerfile/train_vit_gpu.Dockerfile b/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/dockerfile/train_vit_gpu.Dockerfile new file mode 100644 index 000000000..2ddd30753 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/dockerfile/train_vit_gpu.Dockerfile @@ -0,0 +1,149 @@ +# This Dockerfile runs the JAX based Vision transformer training on GPU. +# See https://github.com/google-research/vision_transformer#running-on-cloud +# for more details. +# Here is an example to build this dockerfile: +# PROJECT="your gcp project" +# IMAGE_TAG="trainn_vit_gpu:${USER}-test" +# docker build -f model_oss/jax_vision_transformer/dockerfile/train_vit_gpu.Dockerfile . -t "${IMAGE_TAG}" +# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}" +# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}" + +FROM tensorflow/tensorflow:2.12.0-gpu + +ENV DEBIAN_FRONTEND=noninteractive + +# Install basic libs +RUN apt-get update && apt-get install -y --no-install-recommends \ + cmake \ + curl \ + wget \ + sudo \ + gnupg \ + libsm6 \ + libxext6 \ + libxrender-dev \ + lsb-release \ + ca-certificates \ + build-essential \ + git + +# Copy Apache license. +RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE + +# Get 'vision_transformer' repository from github. +RUN git clone https://github.com/google-research/vision_transformer +# Ser current directory to the downloaded 'vision_transformer' repository. +WORKDIR ./vision_transformer +# Using git reset command to pin it down to a specific version. +RUN git reset --hard e66b4732d44504251197a3da3f5949f3f3ce9ca6 + +# Install required libs +RUN pip install --upgrade pip +# The following pip installs are pinned down versions of those inside +# vit_jax/requirements.txt file. +# NOTE: Using `no-deps` flag to avoid overwriting of +# dependent library versions. For example, +# both `chex` and `jax` can overwrite each others +# `jax-lib` version. +RUN pip install --no-deps absl-py==1.4.0 +RUN pip install --no-deps aqtp==0.0.10 +RUN pip install --no-deps array-record==0.2.0 +RUN pip install --no-deps astunparse==1.6.3 +RUN pip install --no-deps cached-property==1.5.2 +RUN pip install --no-deps cachetools==5.3.0 +RUN pip install --no-deps certifi==2019.11.28 +RUN pip install --no-deps chardet==3.0.4 +RUN pip install --no-deps chex==0.1.7 +RUN pip install --no-deps click==8.1.3 +RUN pip install --no-deps cloudpickle==2.2.1 +RUN pip install --no-deps clu==0.0.9 +RUN pip install --no-deps contextlib2==21.6.0 +RUN pip install --no-deps dacite==1.8.1 +RUN pip install --no-deps dbus-python==1.2.16 +RUN pip install --no-deps decorator==5.1.1 +RUN pip install --no-deps dm-tree==0.1.8 +RUN pip install --no-deps einops==0.6.1 +RUN pip install --no-deps etils==1.3.0 +RUN pip install --no-deps flatbuffers==23.3.3 +RUN pip install --no-deps flax==0.6.10 +RUN pip install --no-deps git+https://github.com/google/flaxformer@9adaa4467cf17703949b9f537c3566b99de1b416 +RUN pip install --no-deps gast==0.4.0 +RUN pip install --no-deps google-auth==2.16.2 +RUN pip install --no-deps google-auth-oauthlib==0.4.6 +RUN pip install --no-deps google-pasta==0.2.0 +RUN pip install --no-deps googleapis-common-protos==1.59.0 +RUN pip install --no-deps grpcio==1.51.3 +RUN pip install --no-deps h5py==3.8.0 +RUN pip install --no-deps idna==2.8 +RUN pip install --no-deps importlib-metadata==6.1.0 +RUN pip install --no-deps importlib-resources==5.12.0 +RUN pip install --no-deps keras==2.12.0 +RUN pip install --no-deps libclang==16.0.0 +RUN pip install --no-deps Markdown==3.4.3 +RUN pip install --no-deps markdown-it-py==2.2.0 +RUN pip install --no-deps MarkupSafe==2.1.2 +RUN pip install --no-deps mdurl==0.1.2 +RUN pip install --no-deps ml-collections==0.1.1 +RUN pip install --no-deps msgpack==1.0.5 +RUN pip install --no-deps nest-asyncio==1.5.6 +RUN pip install --no-deps numpy==1.23.5 +RUN pip install --no-deps oauthlib==3.2.2 +RUN pip install --no-deps opt-einsum==3.3.0 +RUN pip install --no-deps optax==0.1.5 +RUN pip install --no-deps orbax-checkpoint==0.1.6 +RUN pip install --no-deps packaging==23.0 +RUN pip install --no-deps pandas==2.0.1 +RUN pip install --no-deps pip==23.1.2 +RUN pip install --no-deps promise==2.3 +RUN pip install --no-deps protobuf==4.22.1 +RUN pip install --no-deps psutil==5.9.5 +RUN pip install --no-deps pyasn1==0.4.8 +RUN pip install --no-deps pyasn1-modules==0.2.8 +RUN pip install --no-deps Pygments==2.15.1 +RUN pip install --no-deps PyGObject==3.36.0 +RUN pip install --no-deps python-apt==2.0.1+ubuntu0.20.4.1 +RUN pip install --no-deps python-dateutil==2.8.2 +RUN pip install --no-deps pytz==2023.3 +RUN pip install --no-deps PyYAML==6.0 +RUN pip install --no-deps requests==2.22.0 +RUN pip install --no-deps requests-oauthlib==1.3.1 +RUN pip install --no-deps requests-unixsocket==0.2.0 +RUN pip install --no-deps rich==13.3.5 +RUN pip install --no-deps rsa==4.9 +RUN pip install --no-deps scipy==1.10.1 +RUN pip install --no-deps setuptools==67.6.0 +RUN pip install --no-deps six==1.14.0 +RUN pip install --no-deps tensorboard==2.12.0 +RUN pip install --no-deps tensorboard-data-server==0.7.0 +RUN pip install --no-deps tensorboard-plugin-wit==1.8.1 +RUN pip install --no-deps tensorflow==2.12.0 +RUN pip install --no-deps tensorflow-cpu==2.12.0 +RUN pip install --no-deps tensorflow-datasets==4.9.2 +RUN pip install --no-deps tensorflow-estimator==2.12.0 +RUN pip install --no-deps tensorflow-hub==0.13.0 +RUN pip install --no-deps tensorflow-io-gcs-filesystem==0.31.0 +RUN pip install --no-deps tensorflow-metadata==1.13.1 +RUN pip install --no-deps tensorflow-probability==0.20.0 +RUN pip install --no-deps tensorflow-text==2.12.1 +RUN pip install --no-deps tensorstore==0.1.36 +RUN pip install --no-deps termcolor==2.2.0 +RUN pip install --no-deps toml==0.10.2 +RUN pip install --no-deps toolz==0.12.0 +RUN pip install --no-deps tqdm==4.65.0 +RUN pip install --no-deps typing_extensions==4.5.0 +RUN pip install --no-deps tzdata==2023.3 +RUN pip install --no-deps urllib3==1.25.8 +RUN pip install --no-deps Werkzeug==2.2.3 +RUN pip install --no-deps wheel==0.40.0 +RUN pip install --no-deps wrapt==1.14.1 +RUN pip install --no-deps zipp==3.15.0 +# Installing jax at the very end with GPU support. +# NOTE: Not using `no-deps` flag here because +# we need CUDA support. +RUN pip install jax[cuda11_cudnn82]==0.4.6 \ +--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html + +COPY ./model_oss/jax_vision_transformer/vit_config_without_data.py vit_jax/configs/vit.py + +ENV PYTHONPATH ./vit_jax +ENTRYPOINT ["python", "-m", "vit_jax.main"] \ No newline at end of file diff --git a/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/vit_config_without_data.py b/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/vit_config_without_data.py new file mode 100644 index 000000000..242ef8efc --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/vit_config_without_data.py @@ -0,0 +1,27 @@ +"""Returns a config for a Vision Transformer model without asking for data.""" +import ml_collections +from vit_jax.configs import common +from vit_jax.configs import models + + +def get_config(model: str) -> ml_collections.ConfigDict: + """Returns default parameters for finetuning ViT `model`.""" + config = common.get_config() + + get_model_config = getattr(models, f'get_{model}_config') + config.model = get_model_config() + + # These values are often overridden on the command line. + config.base_lr = 0.03 + config.total_steps = 500 + config.warmup_steps = 100 + config.pp = ml_collections.ConfigDict() + config.pp.train = 'train' + config.pp.test = 'test' + config.pp.resize = 448 + config.pp.crop = 384 + + # This value MUST be overridden on the command line. + config.dataset = '' + + return config \ No newline at end of file diff --git a/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/vit_jax2tf.py b/community-content/vertex_vision_model_garden/model_oss/jax_vision_transformer/vit_jax2tf.py new file mode 100644 index 000000000..e69de29bb