diff --git a/mlair/model_modules/abstract_model_class.py b/mlair/model_modules/abstract_model_class.py index 989f4578f78e6566dfca5a63f671ced8120491d8..8af858923207acdc4c471fec0ae7d03f990291b8 100644 --- a/mlair/model_modules/abstract_model_class.py +++ b/mlair/model_modules/abstract_model_class.py @@ -6,6 +6,7 @@ import keras import tensorflow as tf from mlair.helpers import remove_items +from keras import backend as K class AbstractModelClass(ABC): @@ -36,6 +37,7 @@ class AbstractModelClass(ABC): self.__compile_options_is_set = False self._input_shape = input_shape self._output_shape = self.__extract_from_tuple(output_shape) + self.avail_gpus = K.tensorflow_backend._get_available_gpus() def __getattr__(self, name: str) -> Any: """ diff --git a/mlair/model_modules/model_class.py b/mlair/model_modules/model_class.py index b39e97df582d1ee08f6a11b0043638bceec347f0..004fd217411eb0d0e69ce51eb2260a7b120b5b93 100644 --- a/mlair/model_modules/model_class.py +++ b/mlair/model_modules/model_class.py @@ -350,7 +350,7 @@ class MyTowerModel(AbstractModelClass): self.compile_options = {"loss": [keras.losses.mean_squared_error], "metrics": ["mse"]} -class IntelliO3_ts_architecture(AbstractModelClass): +class IntelliO3TsArchitecture(AbstractModelClass): def __init__(self, input_shape: list, output_shape: list): """ @@ -454,7 +454,11 @@ class IntelliO3_ts_architecture(AbstractModelClass): kernel_regularizer=self.regularizer ) - self.model = keras.Model(inputs=X_input, outputs=[out_minor1, out_main]) + model = keras.Model(inputs=X_input, outputs=[out_minor1, out_main]) + if self.avail_gpus == 0: + self.model = model + else: + self.model = keras.utils.multi_gpu_model(model, self.avail_gpus) def set_compile_options(self): self.optimizer = keras.optimizers.SGD(lr=self.initial_lr, momentum=0.9)