131 lines
4.7 KiB
Python
131 lines
4.7 KiB
Python
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
|
|
"""Utilities for Keras classes with v1 and v2 versions."""
|
|
|
|
import tensorflow.compat.v2 as tf
|
|
|
|
from keras.utils.generic_utils import LazyLoader
|
|
|
|
# TODO(b/134426265): Switch back to single-quotes once the issue
|
|
# with copybara is fixed.
|
|
|
|
training = LazyLoader("training", globals(), "keras.engine.training")
|
|
training_v1 = LazyLoader("training_v1", globals(), "keras.engine.training_v1")
|
|
base_layer = LazyLoader("base_layer", globals(), "keras.engine.base_layer")
|
|
base_layer_v1 = LazyLoader(
|
|
"base_layer_v1", globals(), "keras.engine.base_layer_v1"
|
|
)
|
|
callbacks = LazyLoader("callbacks", globals(), "keras.callbacks")
|
|
callbacks_v1 = LazyLoader("callbacks_v1", globals(), "keras.callbacks_v1")
|
|
|
|
|
|
class ModelVersionSelector:
|
|
"""Chooses between Keras v1 and v2 Model class."""
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
use_v2 = should_use_v2()
|
|
cls = swap_class(cls, training.Model, training_v1.Model, use_v2)
|
|
return super(ModelVersionSelector, cls).__new__(cls)
|
|
|
|
|
|
class LayerVersionSelector:
|
|
"""Chooses between Keras v1 and v2 Layer class."""
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
use_v2 = should_use_v2()
|
|
cls = swap_class(cls, base_layer.Layer, base_layer_v1.Layer, use_v2)
|
|
return super(LayerVersionSelector, cls).__new__(cls)
|
|
|
|
|
|
class TensorBoardVersionSelector:
|
|
"""Chooses between Keras v1 and v2 TensorBoard callback class."""
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
use_v2 = should_use_v2()
|
|
start_cls = cls
|
|
cls = swap_class(
|
|
start_cls, callbacks.TensorBoard, callbacks_v1.TensorBoard, use_v2
|
|
)
|
|
if (
|
|
start_cls == callbacks_v1.TensorBoard
|
|
and cls == callbacks.TensorBoard
|
|
):
|
|
# Since the v2 class is not a subclass of the v1 class, __init__ has
|
|
# to be called manually.
|
|
return cls(*args, **kwargs)
|
|
return super(TensorBoardVersionSelector, cls).__new__(cls)
|
|
|
|
|
|
def should_use_v2():
|
|
"""Determine if v1 or v2 version should be used."""
|
|
if tf.executing_eagerly():
|
|
return True
|
|
elif tf.compat.v1.executing_eagerly_outside_functions():
|
|
# Check for a v1 `wrap_function` FuncGraph.
|
|
# Code inside a `wrap_function` is treated like v1 code.
|
|
graph = tf.compat.v1.get_default_graph()
|
|
if getattr(graph, "name", False) and graph.name.startswith(
|
|
"wrapped_function"
|
|
):
|
|
return False
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
|
|
def swap_class(cls, v2_cls, v1_cls, use_v2):
|
|
"""Swaps in v2_cls or v1_cls depending on graph mode."""
|
|
if cls == object:
|
|
return cls
|
|
if cls in (v2_cls, v1_cls):
|
|
return v2_cls if use_v2 else v1_cls
|
|
|
|
# Recursively search superclasses to swap in the right Keras class.
|
|
new_bases = []
|
|
for base in cls.__bases__:
|
|
if (
|
|
use_v2
|
|
and issubclass(base, v1_cls)
|
|
# `v1_cls` often extends `v2_cls`, so it may still call `swap_class`
|
|
# even if it doesn't need to. That being said, it may be the safest
|
|
# not to over optimize this logic for the sake of correctness,
|
|
# especially if we swap v1 & v2 classes that don't extend each
|
|
# other, or when the inheritance order is different.
|
|
or (not use_v2 and issubclass(base, v2_cls))
|
|
):
|
|
new_base = swap_class(base, v2_cls, v1_cls, use_v2)
|
|
else:
|
|
new_base = base
|
|
new_bases.append(new_base)
|
|
cls.__bases__ = tuple(new_bases)
|
|
return cls
|
|
|
|
|
|
def disallow_legacy_graph(cls_name, method_name):
|
|
if not tf.compat.v1.executing_eagerly_outside_functions():
|
|
error_msg = (
|
|
f"Calling `{cls_name}.{method_name}` in graph mode is not "
|
|
f"supported when the `{cls_name}` instance was constructed with "
|
|
f"eager mode enabled. Please construct your `{cls_name}` instance "
|
|
f"in graph mode or call `{cls_name}.{method_name}` with "
|
|
"eager mode enabled."
|
|
)
|
|
raise ValueError(error_msg)
|
|
|
|
|
|
def is_v1_layer_or_model(obj):
|
|
return isinstance(obj, (base_layer_v1.Layer, training_v1.Model))
|