Commit b189793
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -31,40 +31,35 @@ ADD patches/kaggle_datasets.py /root/.local/lib/python3.8/site-packages/kaggle_d | |||
| 31 | 31 | # sitecustomize adds significant latency to ipython kernel startup and should only be added if needed | |
| 32 | 32 | # ADD patches/sitecustomize.py /root/.local/lib/python3.8/site-packages/sitecustomize.py | |
| 33 | 33 | ||
| 34 | - # LIBTPU installed here: | ||
| 35 | - ENV PIP_LIBTPU=/usr/local/lib/python3.8/site-packages/libtpu/libtpu.so | ||
| 36 | - ENV PYTORCH_LIBTPU=/lib/torch-libtpu.so | ||
| 34 | + # Install them all the following together for maximum compatibility. | ||
| 37 | 35 | ||
| 38 | - # Tensorflow wheel: | ||
| 39 | - RUN pip install tensorflow_hub https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-${TENSORFLOW_VERSION}/tensorflow-${TENSORFLOW_VERSION}-cp38-cp38-linux_x86_64.whl \ | ||
| 40 | - tensorflow-addons tensorflow-probability tensorflow-io | ||
| 41 | - RUN curl --output /lib/libtpu.so https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/${TF_LIBTPU_VERSION}/libtpu.so | ||
| 42 | - ADD patches/kaggle_module_resolver.py /usr/local/lib/python3.8/site-packages/tensorflow_hub/kaggle_module_resolver.py | ||
| 43 | - RUN sed -i '/from tensorflow_hub import uncompressed_module_resolver/a from tensorflow_hub import kaggle_module_resolver' /usr/local/lib/python3.8/site-packages/tensorflow_hub/config.py | ||
| 44 | - RUN sed -i '/_install_default_resolvers()/a \ \ registry.resolver.add_implementation(kaggle_module_resolver.KaggleFileResolver())' /usr/local/lib/python3.8/site-packages/tensorflow_hub/config.py | ||
| 36 | + # Install Tensorflow. | ||
| 45 | 37 | ||
| 46 | 38 | # Install Pytorch & related packages | |
| 47 | 39 | # https://cloud.google.com/tpu/docs/pytorch-xla-ug-tpu-vm#changing_pytorch_version | |
| 48 | 40 | # The URL doesn't include patch version. i.e. must use 1.11 instead of 1.11.0 | |
| 49 | 41 | # We need to keep the numpy version the same as the installed tf one but compatible with other installs. | |
| 50 | - RUN NUMPYVER=$(pip freeze | grep -e "^numpy==") && \ | ||
| 51 | - pip install torch==${TORCH_VERSION} https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-${TORCH_VERSION%.*}-cp38-cp38-linux_x86_64.whl torchvision==${TORCHVISION_VERSION} torchtext==${TORCHTEXT_VERSION} torchaudio==${TORCHAUDIO_VERSION} \ | ||
| 52 | - https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-${TENSORFLOW_VERSION}/tensorflow-${TENSORFLOW_VERSION}-cp38-cp38-linux_x86_64.whl \ | ||
| 53 | - "${NUMPYVER}" | ||
| 54 | 42 | ||
| 55 | 43 | # Install JAX & related packages | |
| 56 | 44 | # https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm#install_jax_on_your_cloud_tpu_vm | |
| 57 | - # We need to keep the numpy version the same as the installed tf one but compatible with other installs. | ||
| 58 | - RUN NUMPYVER=$(pip freeze | grep -e "^numpy==") && \ | ||
| 45 | + | ||
| 46 | + # Packages needed by the Notebook editor | ||
| 47 | + | ||
| 48 | + # Additional useful packages should be added here | ||
| 49 | + | ||
| 50 | + RUN pip install tensorflow_hub https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-${TENSORFLOW_VERSION}/tensorflow-${TENSORFLOW_VERSION}-cp38-cp38-linux_x86_64.whl tensorflow-addons tensorflow-probability tensorflow-io \ | ||
| 51 | + pip install torch==${TORCH_VERSION} https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-${TORCH_VERSION%.*}-cp38-cp38-linux_x86_64.whl torchvision==${TORCHVISION_VERSION} torchtext==${TORCHTEXT_VERSION} torchaudio==${TORCHAUDIO_VERSION} \ | ||
| 59 | 52 | pip install jax[tpu]==${JAX_VERSION} -f https://storage.googleapis.com/jax-releases/libtpu_releases.html trax flax optax elegy git+https://github.com/deepmind/dm-haiku jraph distrax \ | |
| 60 | - https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-${TENSORFLOW_VERSION}/tensorflow-${TENSORFLOW_VERSION}-cp38-cp38-linux_x86_64.whl \ | ||
| 61 | - "${NUMPYVER}" | ||
| 53 | + papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1" \ | ||
| 54 | + pandas matplotlib opencv-python-headless librosa scikit-learn | ||
| 62 | 55 | ||
| 63 | - # Packages needed by the Notebook editor: | ||
| 64 | - RUN pip install papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1" | ||
| 56 | + # Tensorflow libtpu: | ||
| 57 | + RUN curl --output /lib/libtpu.so https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/${TF_LIBTPU_VERSION}/libtpu.so | ||
| 65 | 58 | ||
| 66 | - # Additional useful packages should be added here: | ||
| 67 | - RUN pip install pandas matplotlib opencv-python-headless librosa scikit-learn | ||
| 59 | + # Kaggle Model Hub patches: | ||
| 60 | + ADD patches/kaggle_module_resolver.py /usr/local/lib/python3.8/site-packages/tensorflow_hub/kaggle_module_resolver.py | ||
| 61 | + RUN sed -i '/from tensorflow_hub import uncompressed_module_resolver/a from tensorflow_hub import kaggle_module_resolver' /usr/local/lib/python3.8/site-packages/tensorflow_hub/config.py | ||
| 62 | + RUN sed -i '/_install_default_resolvers()/a \ \ registry.resolver.add_implementation(kaggle_module_resolver.KaggleFileResolver())' /usr/local/lib/python3.8/site-packages/tensorflow_hub/config.py | ||
| 68 | 63 | ||
| 69 | 64 | # Monkey-patch the default TPU to the local (TPU VM). | |
| 70 | 65 | RUN sed -i 's/tpu=None,/tpu="local",/' /usr/local/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py | |
0 commit comments