def _init_writer(self, model):
"""Sets file writer."""
if tf.executing_eagerly():
self.writer = tf.summary.create_file_writer(self.log_dir)
if not model.run_eagerly and self.write_graph:
with self.writer.as_default():
tf.summary.graph(K.get_graph())
elif self.write_graph:
self.writer = tf.compat.v1.summary.FileWriter(self.log_dir, K.get_graph())
else:
self.writer = tf.compat.v1.summary.FileWriter(self.log_dir)
def graph_context_for_symbolic_tensors(*args, **kwargs):
"""Returns graph context manager if any of the inputs is a symbolic tensor."""
if any(is_symbolic_tensor(v) for v in list(args) + list(kwargs.values())):
with backend.get_graph().as_default():
yield
else:
yield
def in_keras_graph(self):
# Returns True even if in a subgraph of the Keras graph, such as those
# created by control flow ops.
if tf.executing_eagerly():
return False
return (self._in_keras_graph
or getattr(backend.get_graph(), 'name', None) == 'keras_graph')
def _make_train_function(self):
# Only needed for graph mode and model_to_estimator.
has_recompiled = self._recompile_weights_loss_and_weighted_metrics()
self._check_trainable_weights_consistency()
# If we have re-compiled the loss/weighted metric sub-graphs then create
# train function even if one exists already. This is because
# `_feed_sample_weights` list has been updated on re-compile.
if getattr(self, "train_function", None) is None or has_recompiled:
# Restore the compiled trainable state.
current_trainable_state = self._get_trainable_state()
self._set_trainable_state(self._compiled_trainable_state)
inputs = (self._feed_inputs + self._feed_targets +
self._feed_sample_weights)
if not isinstance(backend.symbolic_learning_phase(), int):
inputs += [backend.symbolic_learning_phase()]
if isinstance(self.optimizer, (list, tuple)):
linear_optimizer = self.optimizer[0]
dnn_optimizer = self.optimizer[1]
else:
linear_optimizer = self.optimizer
dnn_optimizer = self.optimizer
with backend.get_graph().as_default():
with backend.name_scope("training"):
# Training updates
updates = []
linear_updates = linear_optimizer.get_updates(
params=self.linear_model.trainable_weights,
loss=self.total_loss,
)
updates += linear_updates
dnn_updates = dnn_optimizer.get_updates(
params=self.dnn_model.trainable_weights,
loss=self.total_loss,
)
updates += dnn_updates
# Unconditional updates
updates += self.get_updates_for(None)
# Conditional updates relevant to this model
updates += self.get_updates_for(self.inputs)
metrics = self._get_training_eval_metrics()
metrics_tensors = [
m._call_result for m in metrics
if hasattr(m, "_call_result")
]
with backend.name_scope("training"):
# Gets loss and metrics. Updates weights at each call.
fn = backend.function(inputs,
[self.total_loss] + metrics_tensors,
updates=updates,
name="train_function",
**self._function_kwargs)
setattr(self, "train_function", fn)
# Restore the current trainable state
self._set_trainable_state(current_trainable_state)
def clone_model_on_replicas(model, strategy, mode, inputs=None, targets=None):
"""Create a cloned model on each replica."""
with backend.get_graph().as_default(), strategy.scope():
distributed_model = strategy.extended.call_for_each_replica(
_clone_and_build_model, args=(model, mode, inputs, targets))
set_distributed_model(model, mode, distributed_model)
if mode == ModeKeys.TRAIN:
model._make_callback_model(distributed_model)
def _build_distributed_network(model, strategy, mode, inputs=None,
targets=None):
"""Create a cloned model on each replica."""
with K.get_graph().as_default(), strategy.scope():
distributed_model = strategy.extended.call_for_each_replica(
_build_network_on_replica,
args=(model, mode, inputs, targets))
set_distributed_model(model, mode, distributed_model)
def __enter__(self):
call_ctx = self._call_ctx
self._prev_in_call = call_ctx.in_call
self._prev_state = call_ctx._state
call_ctx.in_call = True
call_ctx._state = self._state
# TODO(b/150169018): This logic can be removed after the Functional API
# refactor.
if self._build_graph:
self._prev_in_keras_graph = call_ctx._in_keras_graph
call_ctx._in_keras_graph = (call_ctx._in_keras_graph or getattr(
backend.get_graph(), 'name', None) == 'keras_graph')
def _make_eager_execution_function(model, mode):
"""Makes function to run one step of distributed model eager execution."""
def _per_replica_function(model):
f = model._make_execution_function(mode)
return (f.inputs, f.outputs)
# NOTE(priyag): Try creating a new FuncGraph within DS scope instead of
# using the global one.
strategy = model._distribution_strategy
global_graph = backend.get_graph()
with global_graph.as_default(), strategy.scope():
# First we gather the relevant portions of the model across all
# replicas. `backend._scratch_graph(global_graph)` signals to Keras
# that it should not lift to a separate graph when creating the
# per-replica functions.
with backend._scratch_graph(global_graph):
# Create train ops on each of the devices when we call
# `_per_replica_fit_function`.
grouped = strategy.extended.call_for_each_replica(
_per_replica_function,
args=(get_distributed_model(model, mode), ),
)
grouped_inputs, grouped_outputs = grouped
# Unwrap all the per device values returned from
# `call_for_each_replica`. Unwrapping per device values gives you a
# list of values that can be used to construct a new train function
# that is composed of inputs/outputs on all the devices over which
# the model is distributed.
(all_inputs, all_outputs, _, _) = unwrap_values(
strategy,
grouped_inputs,
grouped_outputs,
with_loss_tensor=(mode != ModeKeys.PREDICT),
)
# Finally, a joint Keras function is created; this one will be created
# in a separate FuncGraph.
return backend.function(
all_inputs,
all_outputs,
name="eager_distributed_{}_function".format(mode),
)
def build_guided_model():
"""Function returning modified model.
Changes gradient function for all ReLu activations
according to Guided Backpropagation.
"""
if "GuidedBackProp" not in ops._gradient_registry._registry:
@ops.RegisterGradient("GuidedBackProp")
def _GuidedBackProp(op, grad):
dtype = op.inputs[0].dtype
return grad * tf.cast(grad > 0., dtype) * \
tf.cast(op.inputs[0] > 0., dtype)
# g = tf.get_default_graph()
g = K.get_graph()
with g.gradient_override_map({'Relu': 'GuidedBackProp'}):
new_model = load_model_first()
return new_model
def _wait_for_variable_initialization(session):
"""Utility to wait for variables to be initialized."""
all_variables = backend._get_variables(backend.get_graph()) # pylint: disable=protected-access
candidate_vars = []
for v in all_variables:
if not getattr(v, '_keras_initialized', False):
candidate_vars.append(v)
if not candidate_vars:
return
while True:
is_initialized = session.run(
[tf.compat.v1.is_variable_initialized(v) for v in candidate_vars])
uninitialized_vars = []
for flag, v in zip(is_initialized, candidate_vars):
if not flag:
uninitialized_vars.append(v)
v._keras_initialized = True # pylint: disable=protected-access
if not uninitialized_vars:
break
""" Functional class for SciANN. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function from tensorflow.python import keras as k from keras import backend as K graph_unique_name = K.get_graph().unique_name from keras.layers import Dense from keras.layers import Concatenate from tensorflow.python.keras.layers import InputSpec from keras.constraints import MinMaxNorm from keras.constraints import Constraint from tensorflow.python.util.tf_export import keras_export from keras.utils import tf_utils from keras import constraints from keras import initializers from keras import regularizers from ..utils import to_list from ..utils import default_constant_initializer from ..utils import default_regularizer from .functional import MLPFunctional from .variable import Variable
def __init__(
self,
input_shape=None,
batch_size=None,
dtype=None,
input_tensor=None,
sparse=None,
name=None,
ragged=None,
type_spec=None,
**kwargs,
):
self._init_input_shape = input_shape
self._init_batch_size = batch_size
self._init_dtype = dtype
self._init_sparse = sparse
self._init_ragged = ragged
self._init_type_spec = type_spec
strategy = tf.distribute.get_strategy()
if (strategy and batch_size is not None
and distributed_training_utils.global_batch_size_supported(
strategy)):
if batch_size % strategy.num_replicas_in_sync != 0:
raise ValueError(
"The `batch_size` argument ({}) must be divisible by "
"the number of replicas ({})".format(
batch_size, strategy.num_replicas_in_sync))
batch_size = batch_size // strategy.num_replicas_in_sync
if "batch_input_shape" in kwargs:
batch_input_shape = kwargs.pop("batch_input_shape")
if input_shape and batch_input_shape:
raise ValueError("Only provide the input_shape OR "
"batch_input_shape argument to "
"InputLayer, not both at the same time.")
# Set the input shape and batch size from the batch_input_shape.
# Note that batch_input_shape can be None (unknown rank) or [] (scalar),
# in which case the batch size must be None.
if batch_input_shape:
batch_size = batch_input_shape[0]
input_shape = batch_input_shape[1:]
if kwargs:
raise ValueError(
f"Unrecognized keyword arguments: {list(kwargs.keys())}")
if sparse and ragged:
raise ValueError(
"Cannot set both sparse and ragged to True in a Keras input.")
if not name:
prefix = "input"
name = prefix + "_" + str(backend.get_uid(prefix))
if not dtype:
if input_tensor is None:
dtype = backend.floatx()
else:
dtype = backend.dtype(input_tensor)
elif input_tensor is not None and input_tensor.dtype != dtype:
raise ValueError(
"`input_tensor.dtype` differs from `dtype`. Received: "
f"input_tensor.dtype={input_tensor.dtype} "
f"but expected dtype={dtype}")
super().__init__(dtype=dtype, name=name)
self.built = True
self.sparse = True if sparse else False
self.ragged = True if ragged else False
self.batch_size = batch_size
self.supports_masking = True
if isinstance(input_shape, tf.TensorShape):
input_shape = tuple(input_shape.as_list())
elif isinstance(input_shape, int):
input_shape = (input_shape, )
if type_spec is not None:
args_that_must_be_none = [
("(input_)shape", self._init_input_shape),
("batch_size", self._init_batch_size),
("dtype", self._init_dtype),
("input_tensor", input_tensor),
("sparse", self._init_sparse),
("ragged", self._init_ragged),
]
for arg_name, arg in args_that_must_be_none:
_assert_other_arg_none(arg_name, arg)
if not tf.compat.v1.executing_eagerly_outside_functions():
raise ValueError(
"Creating Keras inputs from a type_spec is only "
"supported when eager execution is enabled.")
input_tensor = keras_tensor.keras_tensor_from_type_spec(type_spec)
if isinstance(input_tensor, keras_tensor.SparseKerasTensor):
self.sparse = True
if isinstance(input_tensor, keras_tensor.RaggedKerasTensor):
self.ragged = True
self.is_placeholder = True
try:
self._batch_input_shape = tuple(input_tensor.shape.as_list())
except ValueError:
# If the shape cannot be represented as a tuple (e.g. unknown rank)
self._batch_input_shape = None
elif input_tensor is None:
if input_shape is not None:
batch_input_shape = (batch_size, ) + tuple(input_shape)
else:
batch_input_shape = None
graph = backend.get_graph()
with graph.as_default():
input_tensor = backend.placeholder(
shape=batch_input_shape,
dtype=dtype,
name=self.name,
sparse=sparse,
ragged=ragged,
)
self.is_placeholder = True
self._batch_input_shape = batch_input_shape
else:
if tf.compat.v1.executing_eagerly_outside_functions():
if not isinstance(input_tensor, keras_tensor.KerasTensor):
input_tensor = keras_tensor.keras_tensor_from_tensor(
input_tensor)
else:
if not tf_utils.is_symbolic_tensor(input_tensor):
raise ValueError(
"You should not pass an EagerTensor to `Input`. "
"For example, instead of creating an "
"`InputLayer`, you should instantiate your model "
"and directly call it on your input.")
self.is_placeholder = False
try:
self._batch_input_shape = tuple(input_tensor.shape.as_list())
except ValueError:
# If the shape cannot be represented as a tuple (e.g. unknown rank)
self._batch_input_shape = None
# Create an input node.
input_tensor._keras_mask = None
node_module.Node(layer=self, outputs=input_tensor)
# Store type spec
if isinstance(input_tensor, keras_tensor.KerasTensor) or (
tf_utils.is_extension_type(input_tensor)):
self._type_spec = (input_tensor._type_spec) # pylint: disable=protected-access
else:
self._type_spec = tf.TensorSpec(
shape=input_tensor.shape,
dtype=input_tensor.dtype,
name=self.name,
)
def __init__(self,
input_shape=None,
batch_size=None,
dtype=None,
input_tensor=None,
sparse=None,
name=None,
ragged=None,
type_spec=None,
**kwargs):
self._init_input_shape = input_shape
self._init_batch_size = batch_size
self._init_dtype = dtype
self._init_sparse = sparse
self._init_ragged = ragged
self._init_type_spec = type_spec
strategy = tf.distribute.get_strategy()
if strategy and batch_size is not None and \
distributed_training_utils.global_batch_size_supported(strategy):
if batch_size % strategy.num_replicas_in_sync != 0:
raise ValueError(
'The `batch_size` argument ({}) must be divisible by '
'the number of replicas ({})'.format(
batch_size, strategy.num_replicas_in_sync))
batch_size = batch_size // strategy.num_replicas_in_sync
if 'batch_input_shape' in kwargs:
batch_input_shape = kwargs.pop('batch_input_shape')
if input_shape and batch_input_shape:
raise ValueError('Only provide the input_shape OR '
'batch_input_shape argument to '
'InputLayer, not both at the same time.')
batch_size = batch_input_shape[0]
input_shape = batch_input_shape[1:]
if kwargs:
raise ValueError('Unrecognized keyword arguments:', kwargs.keys())
if sparse and ragged:
raise ValueError(
'Cannot set both sparse and ragged to True in a Keras input.')
if not name:
prefix = 'input'
name = prefix + '_' + str(backend.get_uid(prefix))
if not dtype:
if input_tensor is None:
dtype = backend.floatx()
else:
dtype = backend.dtype(input_tensor)
elif input_tensor is not None and input_tensor.dtype != dtype:
raise ValueError(
'`input_tensor.dtype` differs from `dtype`: %s vs. %s' %
(input_tensor.dtype, dtype))
super(InputLayer, self).__init__(dtype=dtype, name=name)
self.built = True
self.sparse = True if sparse else False
self.ragged = True if ragged else False
self.batch_size = batch_size
self.supports_masking = True
if isinstance(input_shape, tf.TensorShape):
input_shape = tuple(input_shape.as_list())
elif isinstance(input_shape, int):
input_shape = (input_shape, )
if type_spec is not None:
args_that_must_be_none = [
('(input_)shape', self._init_input_shape),
('batch_size', self._init_batch_size),
('dtype', self._init_dtype),
('input_tensor', input_tensor),
('sparse', self._init_sparse),
('ragged', self._init_ragged),
]
for arg_name, arg in args_that_must_be_none:
_assert_other_arg_none(arg_name, arg)
if not keras_tensor.keras_tensors_enabled():
raise ValueError(
'Creating Keras inputs from a type_spec is only '
'supported when eager execution is enabled.')
input_tensor = keras_tensor.keras_tensor_from_type_spec(type_spec)
if isinstance(input_tensor, keras_tensor.SparseKerasTensor):
self.sparse = True
if isinstance(input_tensor, keras_tensor.RaggedKerasTensor):
self.ragged = True
self.is_placeholder = True
try:
self._batch_input_shape = tuple(input_tensor.shape.as_list())
except ValueError:
# If the shape cannot be represented as a tuple (e.g. unknown rank)
self._batch_input_shape = None
elif input_tensor is None:
if input_shape is not None:
batch_input_shape = (batch_size, ) + tuple(input_shape)
else:
batch_input_shape = None
graph = backend.get_graph()
with graph.as_default():
input_tensor = backend.placeholder(shape=batch_input_shape,
dtype=dtype,
name=self.name,
sparse=sparse,
ragged=ragged)
self.is_placeholder = True
self._batch_input_shape = batch_input_shape
else:
if keras_tensor.keras_tensors_enabled():
if not isinstance(input_tensor, keras_tensor.KerasTensor):
input_tensor = keras_tensor.keras_tensor_from_tensor(
input_tensor)
else:
if not tf_utils.is_symbolic_tensor(input_tensor):
raise ValueError(
'You should not pass an EagerTensor to `Input`. '
'For example, instead of creating an '
'InputLayer, you should instantiate your model and '
'directly call it on your input.')
self.is_placeholder = False
try:
self._batch_input_shape = tuple(input_tensor.shape.as_list())
except ValueError:
# If the shape cannot be represented as a tuple (e.g. unknown rank)
self._batch_input_shape = None
# Create an input node.
input_tensor._keras_mask = None
node_module.Node(layer=self, outputs=input_tensor)
# Store type spec
if isinstance(input_tensor, keras_tensor.KerasTensor) or (
tf_utils.is_extension_type(input_tensor)):
self._type_spec = input_tensor._type_spec # pylint: disable=protected-access
else:
self._type_spec = tf.TensorSpec(shape=input_tensor.shape,
dtype=input_tensor.dtype,
name=self.name)
