add tf and torch requirements.txt

henrysky · Sep 6, 2024 · f3b5e82 · f3b5e82
1 parent 67f82b6
commit f3b5e82
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 11 deletions.
diff --git a/.github/workflows/ci_tests.yml b/.github/workflows/ci_tests.yml
@@ -10,9 +10,9 @@ jobs:
       matrix:
         include:
           # - { BACKEND: torch, BACKEND_VER: 2.4.0, PAPER_MODELS: true, ALLOW_FAILURE: true }
-          - { BACKEND: torch, BACKEND_VER: 2.4.0, PAPER_MODELS: false, ALLOW_FAILURE: false }
+          - { BACKEND: torch, PAPER_MODELS: false, ALLOW_FAILURE: false }
           # - { BACKEND: tensorflow, BACKEND_VER: 2.17.0, PAPER_MODELS: true, ALLOW_FAILURE: true }
-          - { BACKEND: tensorflow, BACKEND_VER: 2.17.0, PAPER_MODELS: false, ALLOW_FAILURE: false }
+          - { BACKEND: tensorflow, PAPER_MODELS: false, ALLOW_FAILURE: false }
           # - { BACKEND: jax, BACKEND_VER: 0.4.30, PAPER_MODELS: true, ALLOW_FAILURE: true }
           # - { BACKEND: jax, BACKEND_VER: 0.4.30, PAPER_MODELS: false, ALLOW_FAILURE: true }
     env:
@@ -33,13 +33,17 @@ jobs:
           pip install -r requirements.txt
           pip install -r requirements-dev.txt
           # if [ "$(which nvidia-smi)" == "" ] && [ "${BACKEND}" == "torch" ]; then pip install torch~=${{ matrix.BACKEND_VER }} --index-url https://download.pytorch.org/whl/cpu; else pip install ${{ matrix.BACKEND }}~=${{ matrix.BACKEND_VER }}; fi
-          pip install ${{ matrix.BACKEND }}~=${{ matrix.BACKEND_VER }}
+          pip install -r requirements-${{ matrix.BACKEND }}.txt
           pip install .
+      - name: Get Backend version
+        run: echo "BACKEND_VER=$(python -c 'import os; import importlib; importlib.import_module(os.environ["KERAS_BACKEND"]).__version__.split("+")[0]' >> $GITHUB_ENV
+      - name: Get KERAS version
+        run: echo "KERAS_VER=$(python -c 'import keras; keras.__version__' >> $GITHUB_ENV
       - name: Cache test data
         uses: actions/cache@v4
         with:
           path: ${{github.workspace}}/ci_data/
-          key: astronn-${{ runner.os }}-${{ matrix.BACKEND_VER }}-${{ matrix.PAPER_MODELS }}-2024july22
+          key: astronn-${{ runner.os }}-${{ env.BACKEND_VER }}-${{ matrix.PAPER_MODELS }}
       - name: Test astroNN
         continue-on-error: ${{ matrix.ALLOW_FAILURE }}
         if: ${{ matrix.PAPER_MODELS == false}}
@@ -59,11 +63,12 @@ jobs:
         uses: codecov/codecov-action@v4
         env:
           BACKEND: ${{ matrix.BACKEND }}
-          BACKEND_VER: ${{ matrix.BACKEND_VER }}
+          KERAS_VER: ${{ env.KERAS_VER }}
+          BACKEND_VER: ${{ env.BACKEND_VER }}
           PAPER_MODELS: ${{ matrix.PAPER_MODELS }}
         with:
           flags: unittests
-          env_vars: BACKEND, BACKEND_VER, PAPER_MODELS
+          env_vars: BACKEND, KERAS_VER, BACKEND_VER, PAPER_MODELS
           name: codecov-astroNN
           fail_ci_if_error: true
           verbose: true

diff --git a/docs/source/quick_start.rst b/docs/source/quick_start.rst
@@ -36,12 +36,10 @@ Prerequisites
 
 .. literalinclude:: ../../requirements.txt
 
-Currently ``astroNN`` supports both `Tensorflow`_ and `PyTorch`_ as backend. You can install either of them by running:
+Currently ``astroNN`` supports both `Tensorflow`_ and `PyTorch`_ as backend. You can use either one of them. The current required version are:
 
-.. prompt:: bash $
-
-    pip install tensorflow
-    pip install torch
+.. literalinclude:: ../../requirements-torch.txt
+.. literalinclude:: ../../requirements-tensorflow.txt
 
 To plot the model, you will need to install ``graphviz``. On Ubuntu, you can install it by running:
 

diff --git a/requirements-tensorflow.txt b/requirements-tensorflow.txt
@@ -0,0 +1 @@
+tensorflow~=2.17.0
diff --git a/requirements-torch.txt b/requirements-torch.txt
@@ -0,0 +1 @@
+torch~=2.4.0