Proposed changes to #541

CHANGES.rst

@@ -17,7 +17,7 @@ Version 0.0.6   (unreleased)
   ``scico.flax.save_variables`` and ``scico.flax.load_variables``
   respectively.
 • Support ``jaxlib`` and ``jax`` versions 0.4.3 to 0.4.30.
-• Support ``flax`` versions between 0.8.0 and 0.8.3 (inclusive).
+• Support ``flax`` versions 0.8.0 to 0.8.3.
 
 
 

examples/scripts/ct_astra_datagen_foam2.py

@@ -13,8 +13,14 @@
 generated using filtered back projection (FBP).
 """
 
+# isort: off
 import numpy as np
 
+import logging
+import ray
+
+ray.init(logging_level=logging.ERROR)  # need to call init before jax import: ray-project/ray#44087
+
 from scico import plot
 from scico.flax.examples import load_ct_data
 

examples/scripts/ct_astra_modl_train_foam2.py

@@ -40,12 +40,18 @@
 reconstructed images.
 """
 
+# isort: off
 import os
 from functools import partial
 from time import time
 
 import numpy as np
 
+import logging
+import ray
+
+ray.init(logging_level=logging.ERROR)  # need to call init before jax import: ray-project/ray#44087
+
 import jax
 
 from mpl_toolkits.axes_grid1 import make_axes_locatable

examples/scripts/ct_astra_odp_train_foam2.py

@@ -44,12 +44,21 @@
 term. The output of the final stage is the set of reconstructed images.
 """
 
+# isort: off
 import os
 from functools import partial
 from time import time
 
 import numpy as np
 
+import logging
+import ray
+
+ray.init(logging_level=logging.ERROR)  # need to call init before jax import: ray-project/ray#44087
+
+# Set an arbitrary processor count (only applies if GPU is not available).
+os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 import jax
 
 from mpl_toolkits.axes_grid1 import make_axes_locatable
@@ -60,11 +69,7 @@
 from scico.flax.train.traversals import clip_positive, construct_traversal
 from scico.linop.xray.astra import XRayTransform2D
 
-"""
-Prepare parallel processing. Set an arbitrary processor count (only
-applies if GPU is not available).
-"""
-os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 platform = jax.lib.xla_bridge.get_backend().platform
 print("Platform: ", platform)
 

examples/scripts/ct_astra_unet_train_foam2.py

@@ -13,22 +13,29 @@
 by :cite:`jin-2017-unet`.
 """
 
+# isort: off
 import os
 from time import time
 
+import logging
+import ray
+
+ray.init(logging_level=logging.ERROR)  # need to call init before jax import: ray-project/ray#44087
+
+# Set an arbitrary processor count (only applies if GPU is not available).
+os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 import jax
 
+import numpy as np
+
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
 from scico import flax as sflax
 from scico import metric, plot
 from scico.flax.examples import load_ct_data
 
-"""
-Prepare parallel processing. Set an arbitrary processor count (only
-applies if GPU is not available).
-"""
-os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 platform = jax.lib.xla_bridge.get_backend().platform
 print("Platform: ", platform)
 
@@ -190,7 +197,7 @@
     hist = stats_object.history(transpose=True)
     fig, ax = plot.subplots(nrows=1, ncols=2, figsize=(12, 5))
     plot.plot(
-        jax.numpy.vstack((hist.Train_Loss, hist.Eval_Loss)).T,
+        np.vstack((hist.Train_Loss, hist.Eval_Loss)).T,
         x=hist.Epoch,
         ptyp="semilogy",
         title="Loss function",
@@ -201,7 +208,7 @@
         ax=ax[0],
     )
     plot.plot(
-        jax.numpy.vstack((hist.Train_SNR, hist.Eval_SNR)).T,
+        np.vstack((hist.Train_SNR, hist.Eval_SNR)).T,
         x=hist.Epoch,
         title="Metric",
         xlbl="Epoch",

examples/scripts/deconv_datagen_foam1.py

@@ -12,10 +12,17 @@
 training neural network models for deconvolution (deblurring). Foam
 phantoms from xdesign are used to generate the clean images.
 """
+
+# isort: off
 import numpy as np
 
+import logging
+import ray
+
+ray.init(logging_level=logging.ERROR)  # need to call init before jax import: ray-project/ray#44087
+
 from scico import plot
-from scico.flax.examples import load_foam1_blur_data
+from scico.flax.examples import load_blur_data
 
 """
 Read data from cache or generate if not available.
@@ -29,7 +36,7 @@
 nimg = train_nimg + test_nimg
 output_size = 256  # image size
 
-train_ds, test_ds = load_foam1_blur_data(
+train_ds, test_ds = load_blur_data(
     train_nimg,
     test_nimg,
     output_size,

examples/scripts/deconv_modl_train_foam1.py

@@ -41,27 +41,32 @@
 images.
 """
 
+# isort: off
 import os
 from functools import partial
 from time import time
 
 import numpy as np
 
+import logging
+import ray
+
+ray.init(logging_level=logging.ERROR)  # need to call init before jax import: ray-project/ray#44087
+
+# Set an arbitrary processor count (only applies if GPU is not available).
+os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 import jax
 
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
 from scico import flax as sflax
 from scico import metric, plot
-from scico.flax.examples import load_foam1_blur_data
+from scico.flax.examples import load_blur_data
 from scico.flax.train.traversals import clip_positive, construct_traversal
 from scico.linop import CircularConvolve
 
-"""
-Prepare parallel processing. Set an arbitrary processor count (only
-applies if GPU is not available).
-"""
-os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 platform = jax.lib.xla_bridge.get_backend().platform
 print("Platform: ", platform)
 
@@ -87,7 +92,7 @@
 test_nimg = 64  # number of testing images
 nimg = train_nimg + test_nimg
 
-train_ds, test_ds = load_foam1_blur_data(
+train_ds, test_ds = load_blur_data(
     train_nimg,
     test_nimg,
     output_size,

examples/scripts/deconv_odp_train_foam1.py

@@ -49,27 +49,32 @@
 set of deblurred images.
 """
 
+# isort: off
 import os
 from functools import partial
 from time import time
 
 import numpy as np
 
+import logging
+import ray
+
+ray.init(logging_level=logging.ERROR)  # need to call init before jax import: ray-project/ray#44087
+
+# Set an arbitrary processor count (only applies if GPU is not available).
+os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 import jax
 
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
 from scico import flax as sflax
 from scico import metric, plot
-from scico.flax.examples import load_foam1_blur_data
+from scico.flax.examples import load_blur_data
 from scico.flax.train.traversals import clip_positive, construct_traversal
 from scico.linop import CircularConvolve
 
-"""
-Prepare parallel processing. Set an arbitrary processor count (only
-applies if GPU is not available).
-"""
-os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 platform = jax.lib.xla_bridge.get_backend().platform
 print("Platform: ", platform)
 
@@ -95,7 +100,7 @@
 test_nimg = 64  # number of testing images
 nimg = train_nimg + test_nimg
 
-train_ds, test_ds = load_foam1_blur_data(
+train_ds, test_ds = load_blur_data(
     train_nimg,
     test_nimg,
     output_size,

examples/scripts/denoise_dncnn_train_bsds.py

@@ -13,11 +13,15 @@
 with additive Gaussian noise.
 """
 
+# isort: off
 import os
 from time import time
 
 import numpy as np
 
+# Set an arbitrary processor count (only applies if GPU is not available).
+os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 import jax
 
 from mpl_toolkits.axes_grid1 import make_axes_locatable
@@ -26,11 +30,7 @@
 from scico import metric, plot
 from scico.flax.examples import load_image_data
 
-"""
-Prepare parallel processing. Set an arbitrary processor count (only
-applies if GPU is not available).
-"""
-os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 platform = jax.lib.xla_bridge.get_backend().platform
 print("Platform: ", platform)
 

scico/flax/examples/__init__.py

@@ -8,11 +8,11 @@
 """Data utility functions used by Flax example scripts."""
 
 from .data_preprocessing import PaddedCircularConvolve, build_blur_kernel
-from .examples import load_ct_data, load_foam1_blur_data, load_image_data
+from .examples import load_blur_data, load_ct_data, load_image_data
 
 __all__ = [
     "load_ct_data",
-    "load_foam1_blur_data",
+    "load_blur_data",
     "load_image_data",
     "PaddedCircularConvolve",
     "build_blur_kernel",