move keras import to config.py

astroNN/config.py

@@ -3,7 +3,10 @@
 import platform
 import numpy as np
 
-import keras
+try:
+    import keras.src as keras
+except ModuleNotFoundError:
+    import keras
 
 astroNN_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".astroNN")
 _astroNN_MODEL_NAME = "model_weights.keras"  # default astroNN model filename

astroNN/models/__init__.py

@@ -22,7 +22,7 @@
 from astroNN.nn.losses import losses_lookup
 from astroNN.nn.utilities import Normalizer
 from astroNN.shared.dict_tools import dict_list_to_dict_np, list_to_dict
-import keras
+from astroNN.config import keras
 
 __all__ = [
     "load_folder",

astroNN/models/apogee_models.py

@@ -1,7 +1,7 @@
 # ---------------------------------------------------------#
 #   astroNN.models.apogee_models: Contain Apogee Models
 # ---------------------------------------------------------#
-import keras
+from astroNN.config import keras
 import numpy as np
 
 from astroNN.apogee import aspcap_mask

astroNN/models/base_bayesian_cnn.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 from tqdm import tqdm
-import keras
+from astroNN.config import keras
 from astroNN.config import MAGIC_NUMBER, MULTIPROCESS_FLAG
 from astroNN.config import _astroNN_MODEL_NAME
 from astroNN.datasets import H5Loader

astroNN/models/base_cnn.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 from tqdm import tqdm
-import keras
+from astroNN.config import keras
 from astroNN.config import MULTIPROCESS_FLAG
 from astroNN.config import _astroNN_MODEL_NAME
 from astroNN.models.base_master_nn import NeuralNetMaster

astroNN/models/base_master_nn.py

@@ -10,7 +10,7 @@
 
 import numpy as np
 import pylab as plt
-import keras
+from astroNN.config import keras
 
 import astroNN
 from astroNN.config import _astroNN_MODEL_NAME, cpu_gpu_reader

astroNN/models/base_vae.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 from tqdm import tqdm
-import keras
+from astroNN.config import keras
 from astroNN.config import MULTIPROCESS_FLAG
 from astroNN.config import _astroNN_MODEL_NAME
 from astroNN.datasets import H5Loader

astroNN/models/misc_models.py

@@ -1,7 +1,7 @@
 # ---------------------------------------------------------#
 #   astroNN.models.misc_models: Contain Misc. Models
 # ---------------------------------------------------------#
-import keras as tfk
+from astroNN.config import keras as tfk
 
 from astroNN.models.base_bayesian_cnn import BayesianCNNBase
 from astroNN.models.base_cnn import CNNBase

astroNN/nn/callbacks.py

@@ -2,7 +2,7 @@
 import os
 
 import numpy as np
-import keras as tfk
+from astroNN.config import keras as tfk
 
 Callback = tfk.callbacks.Callback
 

astroNN/nn/layers.py

@@ -1,14 +1,5 @@
 import math
-import keras
-try:
-    from keras.layers import InputSpec
-except AttributeError:
-    from keras.layers.input_spec import InputSpec
-
-try:
-    keras_shape_func = keras.backend.shape
-except AttributeError:
-    keras_shape_func = keras.src.backend.shape
+from astroNN.config import keras
 
 epsilon = keras.backend.epsilon
 initializers = keras.initializers
@@ -75,8 +66,8 @@ def __init__(self, name=None, **kwargs):
 
     def call(self, inputs):
         z_mean, z_log_var = inputs
-        batch = keras_shape_func(z_mean)[0]
-        dim = keras_shape_func(z_mean)[1]
+        batch = keras.backend.shape(z_mean)[0]
+        dim = keras.backend.shape(z_mean)[1]
         epsilon = keras.backend.random.normal(shape=(batch, dim))
         return z_mean + keras.backend.numpy.exp(0.5 * z_log_var) * epsilon
 
@@ -108,7 +99,7 @@ def _get_noise_shape(self, inputs):
         if self.noise_shape is None:
             return self.noise_shape
 
-        symbolic_shape = keras_shape_func(inputs)
+        symbolic_shape = keras.backend.shape(inputs)
         noise_shape = [
             symbolic_shape[axis] if shape is None else shape
             for axis, shape in enumerate(self.noise_shape)
@@ -159,10 +150,10 @@ class MCSpatialDropout1D(MCDropout):
     def __init__(self, rate, disable=False, **kwargs):
         super().__init__(rate, disable, **kwargs)
         self.disable_layer = disable
-        self.input_spec = InputSpec(ndim=3)
+        self.input_spec = keras.layers.input_spec.InputSpec(ndim=3)
 
     def _get_noise_shape(self, inputs):
-        input_shape = keras_shape_func(inputs)
+        input_shape = keras.backend.shape(inputs)
         return input_shape[0], 1, input_shape[2]
 
 
@@ -183,10 +174,10 @@ class MCSpatialDropout2D(MCDropout):
     def __init__(self, rate, disable=False, **kwargs):
         super().__init__(rate, disable, **kwargs)
         self.disable_layer = disable
-        self.input_spec = InputSpec(ndim=4)
+        self.input_spec = keras.layers.input_spec.InputSpec(ndim=4)
 
     def _get_noise_shape(self, inputs):
-        input_shape = keras_shape_func(inputs)
+        input_shape = keras.backend.shape(inputs)
         return input_shape[0], 1, 1, input_shape[3]
 
 
@@ -227,7 +218,7 @@ def call(self, inputs, training=None):
             return inputs
         else:
             return inputs * keras.backend.random.normal(
-                shape=keras_shape_func(inputs), mean=1.0, stddev=stddev
+                shape=keras.backend.shape(inputs), mean=1.0, stddev=stddev
             )
 
     def get_config(self):
@@ -531,7 +522,7 @@ def __init__(self, mask, name=None, **kwargs):
         super().__init__(name=name, **kwargs)
 
     def compute_output_shape(self, input_shape):
-        input_shape = keras_shape_func(input_shape)
+        input_shape = keras.backend.shape(input_shape)
         # TODO: convert to keras
         input_shape = input_shape.with_rank_at_least(2)
         return input_shape[:-1].concatenate(self.mask_shape)
@@ -544,7 +535,7 @@ def call(self, inputs, training=None):
         :return: Tensor after applying the layer which is just the masked tensor
         :rtype: tf.Tensor
         """
-        batchsize = keras_shape_func(inputs)[0]
+        batchsize = keras.backend.shape(inputs)[0]
         # need to reshape because tf.keras cannot get the Tensor shape correctly from tf.boolean_mask op
 
         boolean_mask = keras.backend.numpy.any(keras.backend.numpy.not_equal(inputs, self.boolmask), axis=1, keepdims=True)

astroNN/nn/losses.py

@@ -2,7 +2,7 @@
 #   astroNN.nn.losses: losses
 # ---------------------------------------------------------------#
 
-import keras
+from astroNN.config import keras
 
 from astroNN.config import MAGIC_NUMBER
 from astroNN.nn import nn_obj_lookup

astroNN/nn/utilities/generator.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-import keras
+from astroNN.config import keras
 try:
     from keras.trainers.data_adapters.py_dataset_adapter import PyDataset
 except ImportError:

astroNN/shared/nn_tools.py

@@ -3,7 +3,7 @@
 # ---------------------------------------------------------#
 import datetime
 import os
-import keras
+from astroNN.config import keras
 import inspect
 import warnings
 from astroNN.config import _KERAS_BACKEND

tests/custom_model/custom_models.py

@@ -1,17 +1,17 @@
 # ---------------------------------------------------------#
 #   astroNN.models.CIFAR10_CNN: Contain CNN Model
 # ---------------------------------------------------------#
-import keras as tfk
+from astroNN.config import keras
 
 from astroNN.models.base_cnn import CNNBase
 from astroNN.nn.layers import MCDropout
 
-regularizers = tfk.regularizers
-MaxPooling2D, Conv2D, Dense, Flatten, Activation, Input = tfk.layers.MaxPooling2D, tfk.layers.Conv2D, \
-                                                          tfk.layers.Dense, tfk.layers.Flatten, \
-                                                          tfk.layers.Activation, tfk.layers.Input
-max_norm = tfk.constraints.max_norm
-Model = tfk.models.Model
+regularizers = keras.regularizers
+MaxPooling2D, Conv2D, Dense, Flatten, Activation, Input = keras.layers.MaxPooling2D, keras.layers.Conv2D, \
+                                                          keras.layers.Dense, keras.layers.Flatten, \
+                                                          keras.layers.Activation, keras.layers.Input
+max_norm = keras.constraints.max_norm
+Model = keras.models.Model
 
 
 class CustomModel_Test(CNNBase):

tests/test_apogee_model.py

@@ -12,7 +12,7 @@
 from astroNN.nn.callbacks import ErrorOnNaN
 from astroNN.shared.downloader_tools import TqdmUpTo
 
-import keras
+from astroNN.config import keras
 
 _URL_ORIGIN = "https://www.astro.utoronto.ca/~hleung/shared/ci_data/"
 filename = "apogee_dr14_green.h5"

tests/test_layers.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import numpy.testing as npt
-import keras
+from astroNN.config import keras
 
 from astroNN.nn.losses import zeros_loss
 

tests/test_loss_func.py

@@ -1,6 +1,6 @@
 import unittest
 
-import keras
+from astroNN.config import keras
 import numpy as np
 import numpy.testing as npt
 

tests/test_models.py

@@ -5,7 +5,7 @@
 from importlib import import_module
 
 import numpy as np
-import keras
+from astroNN.config import keras
 
 import astroNN
 from astroNN.config import config_path

tests/test_numpy_tools.py

@@ -1,6 +1,6 @@
 import unittest
 
-import keras
+from astroNN.config import keras
 import astropy.units as u
 import numpy as np
 import numpy.testing as npt