ClearML/clearml
1
0
Fork 0
mirror of https://github.com/clearml/clearml synced 2025-06-04 03:47:57 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Extend support for Keras/Tensorflow mix binding

Browse Source
This commit is contained in:
allegroai 2019-06-16 02:28:27 +03:00
parent 3112769ad9
commit e452709492
1 changed files with 36 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
trains/utilities/frameworks.py
View File

@ -1047,7 +1047,8 @@ class PatchTensorFlowEager(object):
class PatchKerasModelIO(object): class PatchKerasModelIO(object):
__main_task = None __main_task = None
__patched = None __patched_keras = None
__patched_tensorflow = None
@staticmethod @staticmethod
def update_current_task(task, **kwargs): def update_current_task(task, **kwargs):
@ -1058,7 +1059,7 @@ class PatchKerasModelIO(object):
@staticmethod @staticmethod
def _patch_model_checkpoint(): def _patch_model_checkpoint():
if 'keras' in sys.modules: if 'keras' in sys.modules and not PatchKerasModelIO.__patched_keras:
try: try:
from keras.engine.network import Network from keras.engine.network import Network
except ImportError: except ImportError:
@ -1071,8 +1072,17 @@ class PatchKerasModelIO(object):
from keras import models as keras_saving from keras import models as keras_saving
except ImportError: except ImportError:
keras_saving = None keras_saving = None
PatchKerasModelIO._patch_io_calls(Network, Sequential, keras_saving) # check that we are not patching anything twice
if 'tensorflow' in sys.modules: if PatchKerasModelIO.__patched_tensorflow:
PatchKerasModelIO.__patched_keras = [
Network if PatchKerasModelIO.__patched_tensorflow[0] != Network else None,
Sequential if PatchKerasModelIO.__patched_tensorflow[1] != Sequential else None,
keras_saving if PatchKerasModelIO.__patched_tensorflow[2] != keras_saving else None,]
else:
PatchKerasModelIO.__patched_keras = [Network, Sequential, keras_saving]
PatchKerasModelIO._patch_io_calls(*PatchKerasModelIO.__patched_keras)
if 'tensorflow' in sys.modules and not PatchKerasModelIO.__patched_tensorflow:
try: try:
# hack: make sure tensorflow.__init__ is called # hack: make sure tensorflow.__init__ is called
import tensorflow import tensorflow
@ -1091,29 +1101,34 @@ class PatchKerasModelIO(object):
from tensorflow.python.keras import models as keras_saving from tensorflow.python.keras import models as keras_saving
except ImportError: except ImportError:
keras_saving = None keras_saving = None
PatchKerasModelIO._patch_io_calls(Network, Sequential, keras_saving)
if PatchKerasModelIO.__patched_keras:
PatchKerasModelIO.__patched_tensorflow = [
Network if PatchKerasModelIO.__patched_keras[0] != Network else None,
Sequential if PatchKerasModelIO.__patched_keras[1] != Sequential else None,
keras_saving if PatchKerasModelIO.__patched_keras[2] != keras_saving else None,]
else:
PatchKerasModelIO.__patched_tensorflow = [Network, Sequential, keras_saving]
PatchKerasModelIO._patch_io_calls(*PatchKerasModelIO.__patched_tensorflow)
@staticmethod @staticmethod
def _patch_io_calls(Network, Sequential, keras_saving): def _patch_io_calls(Network, Sequential, keras_saving):
try: try:
# only patch once if Sequential is not None:
if not PatchKerasModelIO.__patched: Sequential._updated_config = _patched_call(Sequential._updated_config,
PatchKerasModelIO.__patched = True PatchKerasModelIO._updated_config)
if Sequential is not None: Sequential.from_config = _patched_call(Sequential.from_config, PatchKerasModelIO._from_config)
Sequential._updated_config = _patched_call(Sequential._updated_config,
PatchKerasModelIO._updated_config)
Sequential.from_config = _patched_call(Sequential.from_config, PatchKerasModelIO._from_config)
if Network is not None: if Network is not None:
Network._updated_config = _patched_call(Network._updated_config, PatchKerasModelIO._updated_config) Network._updated_config = _patched_call(Network._updated_config, PatchKerasModelIO._updated_config)
Network.from_config = _patched_call(Network.from_config, PatchKerasModelIO._from_config) Network.from_config = _patched_call(Network.from_config, PatchKerasModelIO._from_config)
Network.save = _patched_call(Network.save, PatchKerasModelIO._save) Network.save = _patched_call(Network.save, PatchKerasModelIO._save)
Network.save_weights = _patched_call(Network.save_weights, PatchKerasModelIO._save_weights) Network.save_weights = _patched_call(Network.save_weights, PatchKerasModelIO._save_weights)
Network.load_weights = _patched_call(Network.load_weights, PatchKerasModelIO._load_weights) Network.load_weights = _patched_call(Network.load_weights, PatchKerasModelIO._load_weights)
if keras_saving is not None: if keras_saving is not None:
keras_saving.save_model = _patched_call(keras_saving.save_model, PatchKerasModelIO._save_model) keras_saving.save_model = _patched_call(keras_saving.save_model, PatchKerasModelIO._save_model)
keras_saving.load_model = _patched_call(keras_saving.load_model, PatchKerasModelIO._load_model) keras_saving.load_model = _patched_call(keras_saving.load_model, PatchKerasModelIO._load_model)
except Exception as ex: except Exception as ex:
getLogger(TrainsFrameworkAdapter).warning(str(ex)) getLogger(TrainsFrameworkAdapter).warning(str(ex))