Updating to Keras 3.0 and migrating to PyTorch (#418)

* remove pytest-lazy-fixture as dev dependency and skip test (with WG temp fix) * Test Keras is present (#374) * check if Keras present * change TF to Keras in CI * remove comment * change dependencies in pyproject.toml for Keras 3.0 * Migrate to Keras 3.0 with TF backend (#373) * remove pytest-lazy-fixture as dev dependency and skip test (with WG temp fix) * change tensorflow dependency for cellfinder * replace keras imports from tensorflow to just keras imports * add keras import and reorder * add keras and TF 2.16 to pyproject.toml * comment out TF version check for now * change checkpoint filename for compliance with keras 3. remove use_multiprocessing=False from fit() as it is no longer an input. test_train() passing * add multiprocessing parameters to cube generator constructor and remove from fit() signature (keras3 change) * apply temp garbage collector fix * skip troublesome test * skip running tests on CI on windows * remove commented out TF check * clean commented out code. Explicitly pass use_multiprocessing=False (as before) * remove str conversion before model.save * raise test_detection error for sonarcloud happy * skip running tests on windows on CI * remove filename comment and small edits * Replace TF references in comments and warning messages (#378) * change some old references to TF for the import check * change TF cached model to Keras * Cellfinder with Keras 3.0 and jax backend (#379) * replace tensorflow Tensor with keras tensor * add case for TF prep in prep_model_weights * add different backends to pyproject.toml * add backend configuration to cellfinder init file. tests passing with jax locally * define extra dependencies for cellfinder with different backends. run tox with TF backend * run tox using TF and JAX backend * install TF in brainmapper environment before running tests in CI * add backends check to cellfinder init file * clean up comments * fix tf-nightly import check * specify TF backend in include guard check * clarify comment * remove 'backend' from dependencies specifications * Apply suggestions from code review Co-authored-by: Igor Tatarnikov <[email protected]> --------- Co-authored-by: Igor Tatarnikov <[email protected]> * Run cellfinder with JAX in Windows tests in CI (#382) * use jax backend in brainmapper tests in CI * skip TF backend on windows * fix pip install cellfinder for brainmapper CI tests * add keras env variable for brainmapper CLI tests * fix prep_model_weights * It/keras3 pytorch (#396) * replace tensorflow Tensor with keras tensor * add case for TF prep in prep_model_weights * add different backends to pyproject.toml * add backend configuration to cellfinder init file. tests passing with jax locally * define extra dependencies for cellfinder with different backends. run tox with TF backend * run tox using TF and JAX backend * install TF in brainmapper environment before running tests in CI * add backends check to cellfinder init file * clean up comments * fix tf-nightly import check * specify TF backend in include guard check * clarify comment * remove 'backend' from dependencies specifications * Apply suggestions from code review Co-authored-by: Igor Tatarnikov <[email protected]> * PyTorch runs utilizing multiple cores * PyTorch fix with default models * Tests run on every push for now * Run test on torch backend only * Fixed guard test to set torch as KERAS_BACKEND * KERAS_BACKEND env variable set directly in test_include_guard.yaml * Run test on python 3.11 * Remove tf-nightly from __init__ version check * Added 3.11 to legacy tox config * Changed legacy tox config for real this time * Don't set the wrong max_processing value * Torch is now set as the default backend * Tests only run with torch, updated comments * Unpinned torch version * Add codecov token (#403) * add codecov token * generate xml coverage report * add timeout to testing jobs * Allow turning off classification or detection in GUI (#402) * Allow turning off classification or detection in GUI. * Fix test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactor to fix code analysis errors. * Ensure array is always 2d. * Apply suggestions from code review Co-authored-by: Igor Tatarnikov <[email protected]> --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Igor Tatarnikov <[email protected]> * Support single z-stack tif file for input (#397) * Support single z-stack tif file for input. * Fix commit hook. * Apply review suggestions. * Remove modular asv benchmarks (#406) * remove modular asv benchmarks * recover old structure * remove asv-specific lines from gitignore and manifest * prune benchmarks * Adapt CI so it covers both new and old Macs, and installs required additional dependencies on M1 (#408) * naive attempt at adapting to silicon mac CI * run include guard test on Silicon CI * double-check hdf5 is needed * Optimize cell detection (#398) (#407) * Replace coord map values with numba list/tuple for optim. * Switch to fortran layout for faster update of last dim. * Cache kernel. * jit ball filter. * Put z as first axis to speed z rolling (row-major memory). * Unroll recursion (no perf impact either way). * Parallelize cell cluster splitting. * Parallelize walking for full images. * Cleanup docs and pep8 etc. * Add pre-commit fixes. * Fix parallel always being selected and numba function 1st class warning. * Run hook. * Older python needs Union instead of |. * Accept review suggestion. * Address review changes. * num_threads must be an int. --------- Co-authored-by: Matt Einhorn <[email protected]> * [pre-commit.ci] pre-commit autoupdate (#412) updates: - [github.com/pre-commit/pre-commit-hooks: v4.5.0 → v4.6.0](pre-commit/pre-commit-hooks@v4.5.0...v4.6.0) - [github.com/astral-sh/ruff-pre-commit: v0.3.5 → v0.4.3](astral-sh/ruff-pre-commit@v0.3.5...v0.4.3) - [github.com/psf/black: 24.3.0 → 24.4.2](psf/black@24.3.0...24.4.2) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: sfmig <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Simplify model download (#414) * Simplify model download * Update model cache * Remove jax and tf tests * Standardise the data types for inputs to all be float32 * Force torch to use CPU on arm based macOS during tests * Added PYTORCH_MPS_HIGH_WATERMARK_RATION env variable * Set env variables in test setup * Try to set the default device to cpu in the test itself * Add device call to Conv3D to force cpu * Revert changes, request one cpu left free * Revers the numb cores, don't use arm based mac runner * Merged main, removed torch flags on cellfinder install for guards and brainmapper * Lowercase Torch * Change cache directory --------- Co-authored-by: sfmig <[email protected]> Co-authored-by: Kimberly Meechan <[email protected]> Co-authored-by: Matt Einhorn <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Alessandro Felder <[email protected]> Co-authored-by: Adam Tyson <[email protected]> * Set pooling padding to valid by default on all MaxPooling3D layers * Removed tf error suppression and other tf related functions * Force torch to use cpu device when CELLFINDER_TEST_DEVICE env variable set to cpu * Added nev variable to test step * Use the GITHUB ACTIONS environemntal variable instead * Added docstring for fixture setting device to cpu on arm based mac * Revert changes to no_free_cpus being fixture, and default param * Fixed typo in test_and_deploy.yml * Set multiprocessing to false for the data generators * Update all cache steps to match * Remove reference to TF * Make sure tests can run locally when GITHUB_ACTIONS env variable is missing2 * Removed warning when backend is not configured * Set the label tensor to be float32 to ensure compatibility with mps * Always set KERAS_BACKEND to torch on init * Remove code in __init__ checking for if backend is installed --------- Co-authored-by: sfmig <[email protected]> Co-authored-by: Kimberly Meechan <[email protected]> Co-authored-by: Matt Einhorn <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Alessandro Felder <[email protected]> Co-authored-by: Adam Tyson <[email protected]>
brainglobe · May 28, 2024 · cbdecaf · cbdecaf
1 parent de834d2
commit cbdecaf
Show file tree

Hide file tree

Showing 15 changed files with 157 additions and 177 deletions.
diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml
@@ -37,12 +37,15 @@ jobs:
     name: Run package tests
     timeout-minutes: 60
     runs-on: ${{ matrix.os }}
+    env:
+      KERAS_BACKEND: torch
+      CELLFINDER_TEST_DEVICE: cpu
     strategy:
       matrix:
         # Run all supported Python versions on linux
         os: [ubuntu-latest]
-        python-version: ["3.9", "3.10"]
-        # Include one windows, one macos run each for M1 (latest) and Intel (13)
+        python-version: ["3.9", "3.10", "3.11"]
+        # Include one windows and two macOS (intel based and arm based) runs
         include:
           - os: macos-13
             python-version: "3.10"
@@ -80,11 +83,13 @@ jobs:
        NUMBA_DISABLE_JIT: "1"
 
     steps:
-      - name: Cache tensorflow model
+      - name: Cache brainglobe directory
         uses: actions/cache@v3
         with:
-          path: "~/.cellfinder"
-          key: models-${{ hashFiles('~/.brainglobe/**') }}
+          path: | # ensure we don't cache any interrupted atlas download and extraction, if e.g. we cancel the workflow manually
+            ~/.brainglobe
+            !~/.brainglobe/atlas.tar.gz
+          key: brainglobe
       # Setup pyqt libraries
       - name: Setup qtpy libraries
         uses: tlambert03/setup-qt-libs@v1
@@ -104,13 +109,17 @@ jobs:
     name: Run brainmapper tests to check for breakages
     timeout-minutes: 60
     runs-on: ubuntu-latest
+    env:
+      KERAS_BACKEND: torch
+      CELLFINDER_TEST_DEVICE: cpu
     steps:
-      - name: Cache tensorflow model
+      - name: Cache brainglobe directory
         uses: actions/cache@v3
         with:
-          path: "~/.cellfinder"
-          key: models-${{ hashFiles('~/.brainglobe/**') }}
-
+          path: | # ensure we don't cache any interrupted atlas download and extraction, if e.g. we cancel the workflow manually
+            ~/.brainglobe
+            !~/.brainglobe/atlas.tar.gz
+          key: brainglobe
       - name: Checkout brainglobe-workflows
         uses: actions/checkout@v3
         with:
@@ -124,8 +133,9 @@ jobs:
       - name: Install test dependencies
         run: |
           python -m pip install --upgrade pip wheel
-          # Install latest SHA on this brainglobe-workflows branch
-          python -m pip install git+$GITHUB_SERVER_URL/$GITHUB_REPOSITORY@$GITHUB_SHA
+          # Install cellfinder from the latest SHA on this branch
+          python -m pip install "cellfinder @ git+$GITHUB_SERVER_URL/$GITHUB_REPOSITORY@$GITHUB_SHA"
+
           # Install checked out copy of brainglobe-workflows
           python -m pip install .[dev]
 

diff --git a/.github/workflows/test_include_guard.yaml b/.github/workflows/test_include_guard.yaml
@@ -1,5 +1,5 @@
-name: Test Tensorflow include guards
-# These tests check that the include guards checking for tensorflow's availability
+name: Test Keras include guards
+# These tests check that the include guards checking for Keras availability
 # behave as expected on ubuntu and macOS.
 
 on:
@@ -9,7 +9,7 @@ on:
       - main
 
 jobs:
-  tensorflow_guards:
+  keras_guards:
     name: Test include guards
     strategy:
       matrix:
@@ -24,24 +24,21 @@ jobs:
         with:
           python-version: '3.10'
 
-      - name: Install via pip
-        run: python -m pip install -e .
+      - name: Install cellfinder via pip
+        run: python -m pip install -e "."
 
       - name: Test (working) import
         uses: jannekem/run-python-script-action@v1
+        env:
+          KERAS_BACKEND: torch
         with:
           fail-on-error: true
           script: |
             import cellfinder.core
             import cellfinder.napari
 
-      - name: Uninstall tensorflow-macos on Mac M1
-        if: matrix.os == 'macos-latest'
-        run: python -m pip uninstall -y tensorflow-macos
-
-      - name: Uninstall tensorflow on Ubuntu
-        if: matrix.os == 'ubuntu-latest'
-        run: python -m pip uninstall -y tensorflow
+      - name: Uninstall keras
+        run: python -m pip uninstall -y keras
 
       - name: Test (broken) import
         id: broken_import

diff --git a/cellfinder/__init__.py b/cellfinder/__init__.py
@@ -1,25 +1,31 @@
+import os
 from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
 
+# Check cellfinder is installed
 try:
     __version__ = version("cellfinder")
 except PackageNotFoundError as e:
     raise PackageNotFoundError("cellfinder package not installed") from e
 
-# If tensorflow is not present, tools cannot be used.
+# If Keras is not present, tools cannot be used.
 # Throw an error in this case to prevent invocation of functions.
 try:
-    TF_VERSION = version("tensorflow")
+    KERAS_VERSION = version("keras")
 except PackageNotFoundError as e:
-    try:
-        TF_VERSION = version("tensorflow-macos")
-    except PackageNotFoundError as e:
-        raise PackageNotFoundError(
-            f"cellfinder tools cannot be invoked without tensorflow. "
-            f"Please install tensorflow into your environment to use cellfinder tools. "
-            f"For more information, please see "
-            f"https://github.com/brainglobe/brainglobe-meta#readme."
-        ) from e
+    raise PackageNotFoundError(
+        f"cellfinder tools cannot be invoked without Keras. "
+        f"Please install Keras with a backend into your environment "
+        f"to use cellfinder tools. "
+        f"For more information on Keras backends, please see "
+        f"https://keras.io/getting_started/#installing-keras-3."
+        f"For more information on brainglobe, please see "
+        f"https://github.com/brainglobe/brainglobe-meta#readme."
+    ) from e
+
+
+# Set the Keras backend to torch
+os.environ["KERAS_BACKEND"] = "torch"
 
 __author__ = "Adam Tyson, Christian Niedworok, Charly Rousseau"
 __license__ = "BSD-3-Clause"

diff --git a/cellfinder/core/classify/classify.py b/cellfinder/core/classify/classify.py
@@ -1,10 +1,10 @@
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
+import keras
 import numpy as np
 from brainglobe_utils.cells.cells import Cell
 from brainglobe_utils.general.system import get_num_processes
-from tensorflow import keras
 
 from cellfinder.core import logger, types
 from cellfinder.core.classify.cube_generator import CubeGeneratorFromFile
@@ -48,9 +48,7 @@ def main(
         callbacks = None
 
     # Too many workers doesn't increase speed, and uses huge amounts of RAM
-    workers = get_num_processes(
-        min_free_cpu_cores=n_free_cpus, n_max_processes=max_workers
-    )
+    workers = get_num_processes(min_free_cpu_cores=n_free_cpus)
 
     logger.debug("Initialising cube generator")
     inference_generator = CubeGeneratorFromFile(
@@ -63,6 +61,8 @@ def main(
         cube_width=cube_width,
         cube_height=cube_height,
         cube_depth=cube_depth,
+        use_multiprocessing=False,
+        workers=workers,
     )
 
     model = get_model(
@@ -73,10 +73,9 @@ def main(
     )
 
     logger.info("Running inference")
+    # in Keras 3.0 multiprocessing params are specified in the generator
     predictions = model.predict(
         inference_generator,
-        use_multiprocessing=True,
-        workers=workers,
         verbose=True,
         callbacks=callbacks,
     )

diff --git a/cellfinder/core/classify/cube_generator.py b/cellfinder/core/classify/cube_generator.py
@@ -2,13 +2,13 @@
 from random import shuffle
 from typing import Dict, List, Optional, Tuple, Union
 
+import keras
 import numpy as np
-import tensorflow as tf
 from brainglobe_utils.cells.cells import Cell, group_cells_by_z
 from brainglobe_utils.general.numerical import is_even
+from keras.utils import Sequence
 from scipy.ndimage import zoom
 from skimage.io import imread
-from tensorflow.keras.utils import Sequence
 
 from cellfinder.core import types
 from cellfinder.core.classify.augment import AugmentationParameters, augment
@@ -56,7 +56,14 @@ def __init__(
         translate: Tuple[float, float, float] = (0.05, 0.05, 0.05),
         shuffle: bool = False,
         interpolation_order: int = 2,
+        *args,
+        **kwargs,
     ):
+        # pass any additional arguments not specified in signature to the
+        # constructor of the superclass (e.g.: `use_multiprocessing` or
+        # `workers`)
+        super().__init__(*args, **kwargs)
+
         self.points = points
         self.signal_array = signal_array
         self.background_array = background_array
@@ -218,10 +225,10 @@ def __getitem__(self, index: int) -> Union[
 
         if self.train:
             batch_labels = [cell.type - 1 for cell in cell_batch]
-            batch_labels = tf.keras.utils.to_categorical(
+            batch_labels = keras.utils.to_categorical(
                 batch_labels, num_classes=self.classes
             )
-            return images, batch_labels
+            return images, batch_labels.astype(np.float32)
         elif self.extract:
             batch_info = self.__get_batch_dict(cell_batch)
             return images, batch_info
@@ -252,7 +259,8 @@ def __generate_cubes(
                 (number_images,)
                 + (self.cube_height, self.cube_width, self.cube_depth)
                 + (self.channels,)
-            )
+            ),
+            dtype=np.float32,
         )
 
         for idx, cell in enumerate(cell_batch):
@@ -350,7 +358,14 @@ def __init__(
         translate: Tuple[float, float, float] = (0.2, 0.2, 0.2),
         train: bool = False,  # also return labels
         interpolation_order: int = 2,
+        *args,
+        **kwargs,
     ):
+        # pass any additional arguments not specified in signature to the
+        # constructor of the superclass (e.g.: `use_multiprocessing` or
+        # `workers`)
+        super().__init__(*args, **kwargs)
+
         self.im_shape = shape
         self.batch_size = batch_size
         self.labels = labels
@@ -410,10 +425,10 @@ def __getitem__(self, index: int) -> Union[
 
         if self.train and self.labels is not None:
             batch_labels = [self.labels[k] for k in indexes]
-            batch_labels = tf.keras.utils.to_categorical(
+            batch_labels = keras.utils.to_categorical(
                 batch_labels, num_classes=self.classes
             )
-            return images, batch_labels
+            return images, batch_labels.astype(np.float32)
         else:
             return images
 
@@ -424,7 +439,8 @@ def __generate_cubes(
     ) -> np.ndarray:
         number_images = len(list_signal_tmp)
         images = np.empty(
-            ((number_images,) + self.im_shape + (self.channels,))
+            ((number_images,) + self.im_shape + (self.channels,)),
+            dtype=np.float32,
         )
 
         for idx, signal_im in enumerate(list_signal_tmp):
@@ -433,7 +449,7 @@ def __generate_cubes(
                 images, idx, signal_im, background_im
             )
 
-        return images.astype(np.float16)
+        return images
 
     def __populate_array_with_cubes(
         self,

diff --git a/cellfinder/core/classify/resnet.py b/cellfinder/core/classify/resnet.py
@@ -1,9 +1,11 @@
 from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
 
-from tensorflow import Tensor
-from tensorflow.keras import Model
-from tensorflow.keras.initializers import Initializer
-from tensorflow.keras.layers import (
+from keras import (
+    KerasTensor as Tensor,
+)
+from keras import Model
+from keras.initializers import Initializer
+from keras.layers import (
     Activation,
     Add,
     BatchNormalization,
@@ -14,7 +16,7 @@
     MaxPooling3D,
     ZeroPadding3D,
 )
-from tensorflow.keras.optimizers import Adam, Optimizer
+from keras.optimizers import Adam, Optimizer
 
 #####################################################################
 # Define the types of ResNet
@@ -113,7 +115,7 @@ def non_residual_block(
     activation: str = "relu",
     use_bias: bool = False,
     bn_epsilon: float = 1e-5,
-    pooling_padding: str = "same",
+    pooling_padding: str = "valid",
     axis: int = 3,
 ) -> Tensor:
     """
@@ -131,6 +133,7 @@ def non_residual_block(
     )(x)
     x = BatchNormalization(axis=axis, epsilon=bn_epsilon, name="conv1_bn")(x)
     x = Activation(activation, name="conv1_activation")(x)
+
     x = MaxPooling3D(
         max_pool_size,
         strides=strides,