Merge pull request #1968 from christopherwoodall/update-conda-environ…

…ment-yaml

Update tensorflow conda environment

docker/Dockerfile-base

@@ -10,6 +10,14 @@
 
 FROM nvidia/cuda:11.6.2-cudnn8-runtime-ubuntu20.04
 
+ENV PATH=/opt/conda/bin:$PATH
+# TensorFlow requirement
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/conda/lib/
+
+WORKDIR /armory-repo/
+
+COPY environment.yml /armory-repo/
+
 # Temporary fix for broken nvidia package checksum
 # RUN rm -f /etc/apt/sources.list.d/nvidia-ml.list
 
@@ -33,23 +41,11 @@ RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86
     rm ~/miniconda.sh && \
     /opt/conda/bin/conda clean -tipsy && \
     ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh
-#    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
-#    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
-#    echo "conda activate base" >> ~/.bashrc && \
-#    echo 'alias ll="ls -al"' >> ~/.bashrc
-
-ENV PATH=/opt/conda/bin:$PATH
-
-# TensorFlow requirement
-ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/conda/lib/
 
 # NOTE: using mamba because conda fails when trying to solve for environment
 RUN conda install -c conda-forge -n base mamba \
     && conda clean --all
 
-WORKDIR /armory-repo/
-COPY environment.yml /armory-repo/
-
 RUN mamba env update -f environment.yml -n base --prune \
     && mamba clean --all
 
@@ -60,8 +56,8 @@ RUN mamba env update -f environment.yml -n base --prune \
 # NOTE: Armory requirements and ART requirements are installed here to make patch updates fast and small
 RUN echo "Installing TensorFlow and ART/Armory requirements via pip"
 RUN /opt/conda/bin/pip install --no-cache-dir \
-    tensorflow-datasets==4.6 \
-    tensorflow==2.10 \
+    tensorflow-datasets \
+    tensorflow \
     tensorboardx \
     boto3 \
     opencv-python \

docker/Dockerfile-pytorch-deepspeech

@@ -26,10 +26,6 @@ WORKDIR /workspace
 # ------------------------------------------------------------------
 # DEVELOPER NOTES:
 # ------------------------------------------------------------------
-# TODO: determine if this environment setup is needed
-#  $ ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/cuda/lib64"
-
-# NOTE:
 #  - pytorch-lightning >= 1.5.0 will break Deep Speech 2
 #  - torchmetrics >= 0.8.0 will break pytorch-lightning 1.4
 #  - hydra-lightning installs omegaconf

environment.yml

@@ -5,7 +5,8 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - pytorch=1.12
+  - tensorflow
+  - pytorch
   - torchvision
   - torchaudio
   - cudatoolkit=11.6

pyproject.toml

@@ -77,15 +77,16 @@ engine = [
 
 pytorch = [
     "armory-testbed[engine,datasets,math]",
-    "torch",
+    "torch >=1.12.0,<=1.13.1", # Pytorch 2 will break parts of armory
     "torchaudio",
     "torchvision",
 ]
 
 tensorflow = [
     "armory-testbed[engine,datasets,math]",
     "tf-models-official",
-    "tensorflow >= 2.10.0",
+    "tensorflow == 2.10.0; python_version < '3.10'",
+    "tensorflow >= 2.5.0; python_version >= '3.10'",
 ]
 
 deepspeech = [
@@ -97,13 +98,14 @@ deepspeech = [
     "google-cloud-storage",
     "transformers",
     "pytorch-lightning < 1.5.0",
+    "hydra-core",
 ]
 
 math = [
     "numpy",
     "pandas",
     "scipy >= 1.4.1",
-    "scikit-learn < 1.1.0",   # ART requires scikit-learn >=0.22.2,<1.1.0
+    "scikit-learn>=0.22.2,<1.2.0",   # ART requires scikit-learn >=0.22.2,<1.2.0
     "matplotlib",
 ]