It/keras3 pytorch

.github/workflows/test_and_deploy.yml

@@ -41,7 +41,7 @@ jobs:
       matrix:
         # Run all supported Python versions on linux
         os: [ubuntu-latest]
-        python-version: ["3.9", "3.10"]
+        python-version: ["3.9", "3.10", "3.11"]
         # Include one macos run
         include:
         - os: macos-latest
@@ -98,7 +98,7 @@ jobs:
     name: Run brainmapper tests to check for breakages
     runs-on: ubuntu-latest
     env:
-      KERAS_BACKEND: jax
+      KERAS_BACKEND: torch
     steps:
       - name: Cache Keras model
         uses: actions/cache@v3
@@ -119,8 +119,9 @@ jobs:
       - name: Install test dependencies
         run: |
           python -m pip install --upgrade pip wheel
-          # Install cellfinder from the latest SHA on this branch (Keras with JAX backend)
-          python -m pip install "cellfinder[jax] @ git+$GITHUB_SERVER_URL/$GITHUB_REPOSITORY@$GITHUB_SHA"
+          # Install cellfinder from the latest SHA on this branch (Keras with torch backend)
+          python -m pip install "cellfinder[torch] @ git+$GITHUB_SERVER_URL/$GITHUB_REPOSITORY@$GITHUB_SHA"
+
           # Install checked out copy of brainglobe-workflows
           python -m pip install .[dev]
 

.github/workflows/test_include_guard.yaml

@@ -24,11 +24,13 @@ jobs:
         with:
           python-version: '3.10'
 
-      - name: Install cellfinder via pip, specifying tensorflow as keras' backend
-        run: python -m pip install -e ".[tf]"
+      - name: Install cellfinder via pip, specifying torch as keras' backend
+        run: python -m pip install -e ".[torch]"
 
       - name: Test (working) import
         uses: jannekem/run-python-script-action@v1
+        env:
+          KERAS_BACKEND: torch
         with:
           fail-on-error: true
           script: |

cellfinder/__init__.py

@@ -27,25 +27,22 @@
 # If no backend is configured and installed for Keras, tools cannot be used
 # Check backend is configured
 if not os.getenv("KERAS_BACKEND"):
-    os.environ["KERAS_BACKEND"] = "tensorflow"
-    warnings.warn(
-        "Keras backend not configured, automatically set to Tensorflow"
-    )
+    os.environ["KERAS_BACKEND"] = "torch"
+    warnings.warn("Keras backend not configured, automatically set to Torch")
 
 # Check backend is installed
 if os.getenv("KERAS_BACKEND") in ["tensorflow", "jax", "torch"]:
     backend = os.getenv("KERAS_BACKEND")
     try:
-        backend_package = "tf-nightly" if backend == "tensorflow" else backend
-        BACKEND_VERSION = version(backend_package)
+        BACKEND_VERSION = version(backend)
     except PackageNotFoundError as e:
         raise PackageNotFoundError(
-            f"{backend}, ({backend_package}) set as Keras backend "
+            f"{backend}, ({backend}) set as Keras backend "
             f"but not installed"
         ) from e
 else:
     raise PackageNotFoundError(
-        "Keras backend must be one of 'tensorflow', 'jax', or 'torch'"
+        "Keras backend must be one of 'torch', 'tensorflow', or 'jax'"
     )
 
 

cellfinder/core/classify/classify.py

@@ -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(

cellfinder/core/classify/resnet.py

@@ -1,5 +1,6 @@
 from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
 
+import keras.config
 from keras import (
     KerasTensor as Tensor,
 )
@@ -133,6 +134,10 @@ def non_residual_block(
     )(x)
     x = BatchNormalization(axis=axis, epsilon=bn_epsilon, name="conv1_bn")(x)
     x = Activation(activation, name="conv1_activation")(x)
+
+    if keras.config.backend() == "torch":
+        pooling_padding = "valid"
+
     x = MaxPooling3D(
         max_pool_size,
         strides=strides,

cellfinder/core/train/train_yml.py

@@ -22,7 +22,10 @@
     check_positive_float,
     check_positive_int,
 )
-from brainglobe_utils.general.system import ensure_directory_exists
+from brainglobe_utils.general.system import (
+    ensure_directory_exists,
+    get_num_processes,
+)
 from brainglobe_utils.IO.cells import find_relevant_tiffs
 from brainglobe_utils.IO.yaml import read_yaml_section
 from fancylog import fancylog
@@ -363,6 +366,7 @@ def run(
 
     signal_train, background_train, labels_train = make_lists(tiff_files)
 
+    n_processes = get_num_processes(min_free_cpu_cores=n_free_cpus)
     if test_fraction > 0:
         logger.info("Splitting data into training and validation datasets")
         (
@@ -389,7 +393,8 @@ def run(
             labels=labels_test,
             batch_size=batch_size,
             train=True,
-            use_multiprocessing=False,
+            use_multiprocessing=True,
+            workers=n_processes,
         )
 
         # for saving checkpoints
@@ -408,7 +413,8 @@ def run(
         shuffle=True,
         train=True,
         augment=not no_augment,
-        use_multiprocessing=False,
+        use_multiprocessing=True,
+        workers=n_processes,
     )
     callbacks = []
 

pyproject.toml

@@ -61,14 +61,14 @@ napari = [
 ]
 # Keras backends
 tf = [
-    "tf-nightly==2.16.0.dev20240101", # pinning to same TF as Keras 3.0
+    "tensorflow>=2.16.1",
 ]
 jax = [
     "jax==0.4.20",
     "jaxlib==0.4.20"
 ]
 torch = [
-    "torch==2.1.0"
+    "torch>=2.1.0"
 ]
 
 [project.scripts]
@@ -88,7 +88,7 @@ requires = ["setuptools>=45", "wheel", "setuptools_scm[toml]>=6.2"]
 build-backend = 'setuptools.build_meta'
 
 [tool.black]
-target-version = ['py39', 'py310']
+target-version = ['py39', 'py310','py311']
 skip-string-normalization = false
 line-length = 79
 
@@ -120,13 +120,14 @@ markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"]
 legacy_tox_ini = """
 # For more information about tox, see https://tox.readthedocs.io/en/latest/
 [tox]
-envlist = py{39,310}-{tf,jax}
+envlist = py{39,310,311}-{torch}
 isolated_build = true
 
 [gh-actions]
 python =
-    3.9: py39-{tf,jax}  # On GA python=3.9 job, run tox with the tf and jax environments
-    3.10: py310-{tf,jax}  # On GA python=3.10 job, run tox with the tf and jax environments
+    3.9: py39-{torch}  # On GA python=3.9 job, run tox with the torch environment
+    3.10: py310-{torch}  # On GA python=3.10 job, run tox with the torch environment
+    3.11: py311-{torch}  # On GA python=3.11 job, run tox with the torch environment
 
 [testenv]
 commands = python -m pytest -v --color=yes
@@ -143,9 +144,11 @@ extras =
     napari
     tf: tf
     jax: jax
+    torch: torch
 setenv =
     tf: KERAS_BACKEND = tensorflow
     jax: KERAS_BACKEND = jax
+    torch: KERAS_BACKEND = torch
 passenv =
     NUMBA_DISABLE_JIT
     CI