From c9e0fffd40abe1f3712dd2eef065f9831563fb76 Mon Sep 17 00:00:00 2001 From: Henry Leung Date: Thu, 5 Sep 2024 18:39:43 -0400 Subject: [PATCH] check backend compatability when astroNN is imported --- src/astroNN/__init__.py | 16 ++++++++++++++++ src/astroNN/config.py | 16 ++++------------ src/astroNN/models/base_master_nn.py | 21 ++++++++++----------- src/astroNN/models/base_vae.py | 26 +++++++++++--------------- src/astroNN/nn/layers.py | 6 +++--- 5 files changed, 44 insertions(+), 41 deletions(-) diff --git a/src/astroNN/__init__.py b/src/astroNN/__init__.py index 6c3d76fb..fbb3c630 100644 --- a/src/astroNN/__init__.py +++ b/src/astroNN/__init__.py @@ -4,4 +4,20 @@ from importlib.metadata import version +import keras + version = __version__ = version("astroNN") +_KERAS_BACKEND = keras.backend.backend() + +# check if the backend is tensorflow or pytorch +def check_backend(): + """ + Check if the backend is tensorflow or pytorch + """ + if _KERAS_BACKEND != "tensorflow" and _KERAS_BACKEND != "torch": + raise ImportError( + f"astroNN only supports PyTorch or Tensorflow backend, but your current backend is {keras.backend.backend()}" + ) + + +check_backend() diff --git a/src/astroNN/config.py b/src/astroNN/config.py index 6b100641..53e9fa1f 100644 --- a/src/astroNN/config.py +++ b/src/astroNN/config.py @@ -1,22 +1,12 @@ import configparser import os import platform - -import keras -import numpy as np +from astroNN import _KERAS_BACKEND import importlib +backend_framework = importlib.import_module(_KERAS_BACKEND) astroNN_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".astroNN") _astroNN_MODEL_NAME = "model_weights.keras" # default astroNN model filename -_KERAS_BACKEND = keras.backend.backend() - -supported_backend = ["tensorflow", "torch", "jax"] -if _KERAS_BACKEND not in supported_backend: - raise ImportError( - f"astroNN only support {supported_backend} backend, currently you have '{keras.backend.backend()}' as backend" - ) -else: - backend_framework = importlib.import_module(_KERAS_BACKEND) def config_path(flag=None): @@ -249,3 +239,5 @@ def cpu_gpu_reader(): MULTIPROCESS_FLAG = multiprocessing_flag_reader() ENVVAR_WARN_FLAG = envvar_warning_flag_reader() CUSTOM_MODEL_PATH = custom_model_path_reader() + +__all__ = [_KERAS_BACKEND, MAGIC_NUMBER, MULTIPROCESS_FLAG, ENVVAR_WARN_FLAG, CUSTOM_MODEL_PATH, backend_framework] diff --git a/src/astroNN/models/base_master_nn.py b/src/astroNN/models/base_master_nn.py index 5d30cd7f..0682439c 100644 --- a/src/astroNN/models/base_master_nn.py +++ b/src/astroNN/models/base_master_nn.py @@ -16,7 +16,10 @@ from astroNN.config import _astroNN_MODEL_NAME, cpu_gpu_reader from astroNN.shared.nn_tools import cpu_fallback from astroNN.shared.nn_tools import folder_runnum - +from astroNN.config import ( + _KERAS_BACKEND, + backend_framework, +) epsilon, plot_model = keras.backend.epsilon, keras.utils.plot_model @@ -508,24 +511,20 @@ def jacobian(self, x=None, mean_output=False, mc_num=1, denormalize=False): start_time = time.time() - if keras.backend.backend() == "tensorflow": - import tensorflow as tf - - xtensor = tf.Variable(x_data) + if _KERAS_BACKEND == "tensorflow": + xtensor = backend_framework.Variable(x_data) - with tf.GradientTape(watch_accessed_variables=False) as tape: + with backend_framework.GradientTape(watch_accessed_variables=False) as tape: tape.watch(xtensor) temp = _model(xtensor) if isinstance(temp, dict): temp = temp["output"] jacobian = tape.batch_jacobian(temp, xtensor) - elif keras.backend.backend() == "torch": - import torch - + elif _KERAS_BACKEND == "torch": # add new axis for vmap - xtensor = torch.tensor(x_data, requires_grad=True)[:, None, ...] - jacobian = torch.vmap(torch.func.jacrev(_model), randomness="different")(xtensor) + xtensor = backend_framework.tensor(x_data, requires_grad=True)[:, None, ...] + jacobian = backend_framework.vmap(backend_framework.func.jacrev(_model), randomness="different")(xtensor) else: raise ValueError("Only Tensorflow and PyTorch backend is supported") diff --git a/src/astroNN/models/base_vae.py b/src/astroNN/models/base_vae.py index b7a71de4..fa00bcf6 100644 --- a/src/astroNN/models/base_vae.py +++ b/src/astroNN/models/base_vae.py @@ -289,7 +289,7 @@ def custom_train_step(self, data): y = keras.ops.cast(y["output"], backend_framework.float32) # Run forward pass. - if keras.backend.backend() == "tensorflow": + if _KERAS_BACKEND == "tensorflow": with backend_framework.GradientTape() as tape: encoder_output = self.keras_encoder(x, training=True) if isinstance(encoder_output, dict): @@ -313,7 +313,7 @@ def custom_train_step(self, data): self.keras_model.optimizer.apply_gradients( zip(gradients, self.keras_model.trainable_weights) ) - elif keras.backend.backend() == "torch": + elif _KERAS_BACKEND == "torch": self.keras_model.zero_grad() encoder_output = self.keras_encoder(x, training=True) if isinstance(encoder_output, dict): @@ -936,28 +936,24 @@ def jacobian_latent(self, x=None, mean_output=False, mc_num=1, denormalize=False start_time = time.time() - if keras.backend.backend() == "tensorflow": - import tensorflow as tf + if _KERAS_BACKEND == "tensorflow": + xtensor = backend_framework.Variable(x_data) - xtensor = tf.Variable(x_data) - - with tf.GradientTape(watch_accessed_variables=False) as tape: + with backend_framework.GradientTape(watch_accessed_variables=False) as tape: tape.watch(xtensor) temp = _model(xtensor) - jacobian = tf.squeeze(tape.batch_jacobian(temp, xtensor)) - elif keras.backend.backend() == "torch": - import torch - - xtensor = torch.tensor(x_data, requires_grad=True) - jacobian = torch.autograd.functional.jacobian(_model, xtensor) + jacobian = backend_framework.squeeze(tape.batch_jacobian(temp, xtensor)) + elif _KERAS_BACKEND == "torch": + xtensor = backend_framework.tensor(x_data, requires_grad=True) + jacobian = backend_framework.autograd.functional.jacobian(_model, xtensor) else: raise ValueError("Only Tensorflow and PyTorch backend is supported") - jacobian = tf.squeeze(tape.batch_jacobian(temp, xtensor)) + jacobian = backend_framework.squeeze(tape.batch_jacobian(temp, xtensor)) if mean_output is True: - jacobian_master = tf.reduce_mean(jacobian, axis=0).numpy() + jacobian_master = backend_framework.reduce_mean(jacobian, axis=0).numpy() else: jacobian_master = jacobian.numpy() diff --git a/src/astroNN/nn/layers.py b/src/astroNN/nn/layers.py index fc120bd6..9a70f17e 100644 --- a/src/astroNN/nn/layers.py +++ b/src/astroNN/nn/layers.py @@ -1,7 +1,7 @@ import math import keras -from astroNN.config import backend_framework +from astroNN.config import _KERAS_BACKEND, backend_framework class KLDivergenceLayer(keras.layers.Layer): """ @@ -337,12 +337,12 @@ def loop_fn(i): return self.layer(inputs) # vectorizing operation depends on backend - if keras.backend.backend() == "torch": + if _KERAS_BACKEND == "torch": outputs = backend_framework.vmap( loop_fn, randomness="different", in_dims=0 )(self.arange_n) # TODO: tensorflow vectorized_map traced operation so there is no randomness which affects e.g., dropout - # elif keras.backend.backend() == "tensorflow": + # elif _KERAS_BACKEND == "tensorflow": # outputs = backend_framework.vectorized_map(loop_fn, self.arange_n) else: # fallback to simple for loop outputs = [self.layer(inputs) for _ in range(self.n)]