def _init_writer(self, model):
"""Sets file writer."""
if context.executing_eagerly():
self.writer = summary_ops_v2.create_file_writer(self.log_dir)
if not model.run_eagerly and self.write_graph:
with self.writer.as_default():
summary_ops_v2.graph(K.get_graph())
elif self.write_graph:
self.writer = tf_summary.FileWriter(self.log_dir, K.get_graph())
else:
self.writer = tf_summary.FileWriter(self.log_dir)
def _make_eager_execution_function(model, mode):
"""Makes function to run one step of distributed model eager execution."""
strategy = model._distribution_strategy
if not model._grouped_model:
clone_model_on_replicas(
model, strategy, make_callback_model=(mode == 'train'))
def _per_device_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.
with K.get_graph().as_default(), strategy.scope():
# Create train ops on each of the devices when we call
# `_per_device_fit_function`.
(grouped_inputs, grouped_outputs) = strategy.call_for_each_replica(
_per_device_function, args=(model._grouped_model,))
# 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 inptus/outputs
# on all the devices over which the model is distributed.
(all_inputs, all_outputs, _, _) = distributed_training_utils.unwrap_values(
strategy,
grouped_inputs,
grouped_outputs,
with_loss_tensor=(mode != 'predict'))
return K.function(
all_inputs,
all_outputs,
name='eager_distributed_{}_function'.format(mode))
def get_gradients(self, loss, params):
"""Returns gradients of `loss` with respect to `params`.
Arguments:
loss: Loss tensor.
params: List of variables.
Returns:
List of gradient tensors.
Raises:
ValueError: In case any gradient cannot be computed (e.g. if gradient
function not implemented).
"""
params = nest.flatten(params)
with backend.get_graph().as_default():
grads = gradients.gradients(loss, params)
for grad, param in zip(grads, params):
if grad is None:
raise ValueError("Variable {} has `None` for gradient. "
"Please make sure that all of your ops have a "
"gradient defined (i.e. are differentiable). "
"Common ops without gradient: "
"K.argmax, K.round, K.eval.".format(param))
if hasattr(self, "clipnorm"):
grads = [clip_ops.clip_by_norm(g, self.clipnorm) for g in grads]
if hasattr(self, "clipvalue"):
grads = [
clip_ops.clip_by_value(g, -self.clipvalue, self.clipvalue)
for g in grads
]
return grads
def clone_model_on_replicas(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(
_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 is_in_keras_graph():
"""Returns if currently executing inside of a Keras graph."""
# Returns True even if in a subgraph of the Keras graph, such as those
# created by control flow ops.
if context.executing_eagerly():
return False
return (getattr(backend.get_graph(), 'name', None) == 'keras_graph' or
getattr(_call_context, 'in_keras_graph', False))
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 _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(K.symbolic_learning_phase(), int):
inputs += [K.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 K.get_graph().as_default():
with K.name_scope('training'):
# Training updates
updates = []
linear_updates = linear_optimizer.get_updates(
params=self.linear_model.trainable_weights, # pylint: disable=protected-access
loss=self.total_loss)
updates += linear_updates
dnn_updates = dnn_optimizer.get_updates(
params=self.dnn_model.trainable_weights, # pylint: disable=protected-access
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') # pylint: disable=protected-access
]
with K.name_scope('training'):
# Gets loss and metrics. Updates weights at each call.
fn = K.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 __call__(self, fpn_features, boxes, outer_boxes, classes, is_training):
"""Generate the detection priors from the box detections and FPN features.
This corresponds to the Fig. 4 of the ShapeMask paper at
https://arxiv.org/pdf/1904.03239.pdf
Args:
fpn_features: a dictionary of FPN features.
boxes: a float tensor of shape [batch_size, num_instances, 4]
representing the tight gt boxes from dataloader/detection.
outer_boxes: a float tensor of shape [batch_size, num_instances, 4]
representing the loose gt boxes from dataloader/detection.
classes: a int Tensor of shape [batch_size, num_instances]
of instance classes.
is_training: training mode or not.
Returns:
instance_features: a float Tensor of shape [batch_size * num_instances,
mask_crop_size, mask_crop_size, num_downsample_channels]. This is the
instance feature crop.
detection_priors: A float Tensor of shape [batch_size * num_instances,
mask_size, mask_size, 1].
"""
with backend.get_graph().as_default(), tf.name_scope('prior_mask'):
batch_size, num_instances, _ = boxes.get_shape().as_list()
outer_boxes = tf.cast(outer_boxes, tf.float32)
boxes = tf.cast(boxes, tf.float32)
instance_features = spatial_transform_ops.multilevel_crop_and_resize(
fpn_features, outer_boxes, output_size=self._mask_crop_size)
instance_features = self._shape_prior_fc(instance_features)
shape_priors = self._get_priors()
# Get uniform priors for each outer box.
uniform_priors = tf.ones([
batch_size, num_instances, self._mask_crop_size,
self._mask_crop_size
])
uniform_priors = spatial_transform_ops.crop_mask_in_target_box(
uniform_priors, boxes, outer_boxes, self._mask_crop_size)
# Classify shape priors using uniform priors + instance features.
prior_distribution = self._classify_shape_priors(
tf.cast(instance_features, tf.float32), uniform_priors,
classes)
instance_priors = tf.gather(shape_priors, classes)
instance_priors *= tf.expand_dims(tf.expand_dims(tf.cast(
prior_distribution, tf.float32),
axis=-1),
axis=-1)
instance_priors = tf.reduce_sum(instance_priors, axis=2)
detection_priors = spatial_transform_ops.crop_mask_in_target_box(
instance_priors, boxes, outer_boxes, self._mask_crop_size)
return instance_features, detection_priors
def _build_distributed_network(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(
_build_network_on_replica, args=(model, mode, inputs, targets))
set_distributed_model(model, mode, distributed_model)
def write_epoch_models(self, mode: str) -> None:
with self.tf_summary_writers[mode].as_default(
), summary_ops_v2.always_record_summaries():
summary_ops_v2.graph(backend.get_graph(), step=0)
for model in self.network.epoch_models:
summary_writable = (model.__class__.__name__ == 'Sequential'
or (hasattr(model, '_is_graph_network')
and model._is_graph_network))
if summary_writable:
summary_ops_v2.keras_model(model.model_name, model, step=0)
def begin(self):
"""Called once before using the session.
When called, the default graph is the one that will be launched in the
session. The hook can modify the graph by adding new operations to it.
After the `begin()` call the graph will be finalized and the other callbacks
can not modify the graph anymore. Second call of `begin()` on the same
graph, should not change the graph.
"""
from tensorflow.python.keras.backend import get_graph
_graph = get_graph()
try:
self._batch_size_tensor = None
for tensor in _graph.as_graph_def().node:
if "IteratorGetNext" in tensor.name:
_input_tensor = _graph.get_tensor_by_name(tensor.name +
":0")
try:
self._batch_size_tensor = tf.shape(
input=_input_tensor)[0]
except TypeError: # Ragged Tensor
self._batch_size_tensor = _input_tensor.bounding_shape(
)[0]
break
else:
raise RuntimeError(
"Tensor `{}` could not be found. "
"Make sure you are using tf.data API".format(
"IteratorGetNext"))
except RuntimeError:
raise
except Exception as e:
raise RuntimeError(
"Impossible to fetch the tensor: `IteratorGetNext`. Make sure you are using tf.data API."
) from e
self._global_step_tensor = tf.compat.v1.train.get_or_create_global_step(
)
try:
self._AMP_loss_scale_tensor = _graph.get_tensor_by_name(
"current_loss_scale/Read/ReadVariableOp:0")
self._AMP_steps_since_last_loss_scale = _graph.get_tensor_by_name(
"current_loss_scale/Read/ReadVariableOp:0")
except RuntimeError:
raise
# TF-AMP is not activated
except Exception:
pass
def build_model(self, params, mode=None):
if self._keras_model is None:
with backend.get_graph().as_default():
outputs = self.model_outputs(self._input_layer, mode)
model = tf.keras.models.Model(
inputs=self._input_layer, outputs=outputs, name='retinanet')
assert model is not None, 'Fail to build tf.keras.Model.'
model.optimizer = self.build_optimizer()
self._keras_model = model
return self._keras_model
def modify_backprop(model, name):
g = K.get_graph()
with g.gradient_override_map({'ReLU': name}):
# get layers that have an activation
layer_dict = [layer for layer in model.layers[1:] if hasattr(layer, 'activation')]
# replace elu activation
for layer in layer_dict:
if layer.activation == tf.keras.activations.relu:
layer.activation = tf.nn.relu
# re-instanciate a new model
new_model = model
return new_model
def __call__(self, features, is_training=None):
scores_outputs = {}
box_outputs = {}
with backend.get_graph().as_default(), tf.name_scope('rpn_head'):
for level in range(self._min_level, self._max_level + 1):
scores_output, box_output = self._shared_rpn_heads(
features[level], self._anchors_per_location, level, is_training)
scores_outputs[level] = scores_output
box_outputs[level] = box_output
return scores_outputs, box_outputs
def write_graph(self, model: k.Model):
"""Sets Keras model and writes graph if specified."""
if model and self.is_write_graph:
with self.writer.as_default(), summary_ops_v2.always_record_summaries():
if not model.run_eagerly:
summary_ops_v2.graph(get_graph(), step=0)
summary_writable = (
model._is_graph_network or # pylint: disable=protected-access
model.__class__.__name__ == 'Sequential') # pylint: disable=protected-access
if summary_writable:
summary_ops_v2.keras_model('keras', model, step=0)
def compute_model_statistics(batch_size, is_training=True):
"""Compute number of parameters and FLOPS."""
options = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
options['output'] = 'none'
from tensorflow.python.keras.backend import get_graph
flops = tf.compat.v1.profiler.profile(get_graph(),
options=options).total_float_ops
flops_per_image = flops / batch_size
logging.info(
'[%s Compute Statistics] %.1f GFLOPS/image' %
("Training" if is_training else "Inference", flops_per_image / 1e9))
def on_begin(self, state):
if self.write_graph:
with self.summary_writers['train'].as_default():
with summary_ops_v2.always_record_summaries():
summary_ops_v2.graph(backend.get_graph(), step=0)
for name, model in self.network.model.items():
summary_writable = (model._is_graph_network
or model.__class__.__name__
== 'Sequential')
if summary_writable:
summary_ops_v2.keras_model(name, model, step=0)
if self.embeddings_freq:
self._configure_embeddings()
def __call__(self, fpn_features, is_training=None):
"""Returns outputs of RetinaNet head."""
class_outputs = {}
box_outputs = {}
with backend.get_graph().as_default(), tf.name_scope('retinanet'):
for level in range(self._min_level, self._max_level + 1):
features = fpn_features[level]
class_outputs[level] = self.class_net(
features, level, is_training=is_training)
box_outputs[level] = self.box_net(
features, level, is_training=is_training)
return class_outputs, box_outputs
def __call__(self, multilevel_features, is_training=None):
"""Returns the FPN features for a given multilevel features.
Args:
multilevel_features: a `dict` containing `int` keys for continuous feature
levels, e.g., [2, 3, 4, 5]. The values are corresponding features with
shape [batch_size, height_l, width_l, num_filters].
is_training: `bool` if True, the model is in training mode.
Returns:
a `dict` containing `int` keys for continuous feature levels
[min_level, min_level + 1, ..., max_level]. The values are corresponding
FPN features with shape [batch_size, height_l, width_l, fpn_feat_dims].
"""
input_levels = list(multilevel_features.keys())
if min(input_levels) > self._min_level:
raise ValueError('The minimum backbone level %d should be ' %
(min(input_levels)) +
'less or equal to FPN minimum level %d.:' %
(self._min_level))
backbone_max_level = min(max(input_levels), self._max_level)
with backend.get_graph().as_default(), tf.name_scope('fpn'):
# Adds lateral connections.
feats_lateral = {}
for level in range(self._min_level, backbone_max_level + 1):
feats_lateral[level] = self._lateral_conv2d_op[level](
multilevel_features[level])
# Adds top-down path.
feats = {backbone_max_level: feats_lateral[backbone_max_level]}
for level in range(backbone_max_level - 1, self._min_level - 1,
-1):
feats[level] = spatial_transform_ops.nearest_upsampling(
feats[level + 1], 2) + feats_lateral[level]
# Adds post-hoc 3x3 convolution kernel.
for level in range(self._min_level, backbone_max_level + 1):
feats[level] = self._post_hoc_conv2d_op[level](feats[level])
# Adds coarser FPN levels introduced for RetinaNet.
for level in range(backbone_max_level + 1, self._max_level + 1):
feats_in = feats[level - 1]
if level > backbone_max_level + 1:
feats_in = self._activation_op(feats_in)
feats[level] = self._coarse_conv2d_op[level](feats_in)
if self._use_batch_norm:
# Adds batch_norm layer.
for level in range(self._min_level, self._max_level + 1):
feats[level] = self._norm_activations[level](
feats[level], is_training=is_training)
return feats
def __call__(self,
prior_conditioned_features,
class_probs,
is_training=None):
"""Generate instance masks from FPN features and detection priors.
This corresponds to the Fig. 5-6 of the ShapeMask paper at
https://arxiv.org/pdf/1904.03239.pdf
Args:
prior_conditioned_features: a float Tensor of shape [batch_size *
num_instances, mask_crop_size, mask_crop_size, num_downsample_channels].
This is the instance feature crop.
class_probs: a float Tensor of shape [batch_size * num_instances,
mask_crop_size, mask_crop_size, 1]. This is the class probability of
instance segmentation.
is_training: a bool indicating whether in training mode.
Returns:
mask_outputs: instance mask prediction as a float Tensor of shape
[batch_size * num_instances, mask_size, mask_size, num_classes].
"""
with backend.get_graph().as_default(), tf.name_scope(
'affinity-masknet'):
# Extract the foreground mean features
point_samp_prob_thr = 1. / (1. + tf.exp(-self._coarse_mask_thr))
point_samp_prob_thr = tf.cast(point_samp_prob_thr,
class_probs.dtype)
class_probs = tf.where(
tf.greater(class_probs, point_samp_prob_thr), class_probs,
tf.zeros_like(class_probs))
weighted_features = class_probs * prior_conditioned_features
sum_class_vector = tf.reduce_sum(input_tensor=class_probs,
axis=(1, 2)) + tf.constant(
1e-20, class_probs.dtype)
instance_embedding = tf.reduce_sum(input_tensor=weighted_features,
axis=(1, 2)) / sum_class_vector
# Take the difference between crop features and mean instance features.
instance_features = prior_conditioned_features - tf.reshape(
instance_embedding, (-1, 1, 1, self._num_downsample_channels))
# Decoder to generate upsampled segmentation mask.
affinity_output_features = self.finemask_decoder_net(
instance_features, is_training)
# Predict per-class instance masks.
affinity_mask_classes = self._class_predict_conv(
affinity_output_features)
return affinity_mask_classes
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 __call__(self, features, mask_logits, classes, is_training):
"""Generate instance masks from FPN features and detection priors.
This corresponds to the Fig. 5-6 of the ShapeMask paper at
https://arxiv.org/pdf/1904.03239.pdf
Args:
features: a float Tensor of shape
[batch_size, num_instances, mask_crop_size, mask_crop_size,
num_downsample_channels]. This is the instance feature crop.
mask_logits: a float Tensor of shape
[batch_size, num_instances, mask_crop_size, mask_crop_size] indicating
predicted mask logits.
classes: a int Tensor of shape [batch_size, num_instances]
of instance classes.
is_training: a bool indicating whether in training mode.
Returns:
mask_outputs: instance mask prediction as a float Tensor of shape
[batch_size, num_instances, mask_size, mask_size].
"""
# Extract the foreground mean features
# with tf.variable_scope('fine_mask', reuse=tf.AUTO_REUSE):
with backend.get_graph().as_default(), tf.name_scope('fine_mask'):
mask_probs = tf.nn.sigmoid(mask_logits)
# Compute instance embedding for hard average.
binary_mask = tf.cast(tf.greater(mask_probs, 0.5), features.dtype)
instance_embedding = tf.reduce_sum(
features * tf.expand_dims(binary_mask, axis=-1), axis=(2, 3))
instance_embedding /= tf.expand_dims(
tf.reduce_sum(binary_mask, axis=(2, 3)) + 1e-20, axis=-1)
# Take the difference between crop features and mean instance features.
features -= tf.expand_dims(tf.expand_dims(instance_embedding,
axis=2),
axis=2)
features += self._fine_mask_fc(tf.expand_dims(mask_probs, axis=-1))
# Decoder to generate upsampled segmentation mask.
mask_logits = self.decoder_net(features, is_training)
if self._use_category_for_mask:
mask_logits = tf.transpose(mask_logits, [0, 1, 4, 2, 3])
mask_logits = tf.gather(mask_logits,
tf.expand_dims(classes, -1),
batch_dims=2)
mask_logits = tf.squeeze(mask_logits, axis=2)
else:
mask_logits = mask_logits[..., 0]
return mask_logits
def _build_distributed_network(model, strategy, inputs=None, targets=None,
mode=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, inputs, targets, mode))
if mode is ModeKeys.TRAIN:
model._distributed_model_train = distributed_model
elif mode is ModeKeys.TEST:
model._distributed_model_test = distributed_model
elif mode is ModeKeys.PREDICT:
model._distributed_model_predict = distributed_model
else:
model._distributed_model = distributed_model
def clone_model_on_replicas(model, strategy, make_callback_model=False,
inputs=None, targets=None, mode=None):
"""Create a cloned model on each replica."""
with K.get_graph().as_default(), strategy.scope():
grouped_model = strategy.extended.call_for_each_replica(
_clone_and_build_model, args=(model, inputs, targets, mode))
if mode is _Mode.TRAIN:
model._grouped_model_train = grouped_model
elif mode is _Mode.TEST:
model._grouped_model_test = grouped_model
elif mode is _Mode.PREDICT:
model._grouped_model_predict = grouped_model
else:
model._grouped_model = grouped_model
if make_callback_model:
model._make_callback_model(grouped_model)
def clone_model_on_replicas(model, strategy, make_callback_model=False,
inputs=None, targets=None, mode=None):
"""Create a cloned model on each replica."""
with K.get_graph().as_default(), strategy.scope():
distributed_model = strategy.extended.call_for_each_replica(
_clone_and_build_model, args=(model, inputs, targets, mode))
if mode is ModeKeys.TRAIN:
model._distributed_model_train = distributed_model
elif mode is ModeKeys.TEST:
model._distributed_model_test = distributed_model
elif mode is ModeKeys.PREDICT:
model._distributed_model_predict = distributed_model
else:
model._distributed_model = distributed_model
if make_callback_model:
model._make_callback_model(distributed_model)
def _get_graph(explainer):
""" Common utility to get the graph for the tensorflow-based explainer.
Parameters
----------
explainer : Explainer
One of the tensorflow-based explainers.
"""
_import_tf()
if not tf.executing_eagerly():
return explainer.session.graph
else:
from tensorflow.python.keras import backend
graph = backend.get_graph()
return graph
def broadcast_global_variables(root_rank):
"""Used to broadcast global variables"""
variables = backend._get_variables(backend.get_graph())
candidate_vars = []
for v in variables:
if getattr(v, "_keras_initialized", False):
candidate_vars.append(v)
op_list = []
if candidate_vars:
for var in candidate_vars:
inputs = [var]
outputs = hccl_ops.broadcast(tensor=inputs, root_rank=root_rank)
if outputs is not None:
op_list.append(outputs[0].op)
op_list.append(state_ops.assign(var, outputs[0]))
return control_flow_ops.group(op_list)
def call_context(layer):
"""Scope that marks when we are currently inside a Layer/Model's `call`."""
was_in_call = is_in_call_context()
was_frozen = is_in_frozen_context()
was_in_keras_graph = getattr(_call_context, 'in_keras_graph', False)
_call_context.in_call = True
_call_context.in_keras_graph = (was_in_keras_graph or getattr(
backend.get_graph(), 'name', None) == 'keras_graph')
if not layer.trainable:
_call_context.frozen = True
try:
yield
finally:
_call_context.in_call = was_in_call
_call_context.frozen = was_frozen
_call_context.in_keras_graph = was_in_keras_graph
def __call__(self, inputs, is_training=None):
"""Returns the ResNet model for a given size and number of output classes.
Args:
inputs: a `Tesnor` with shape [batch_size, height, width, 3] representing
a batch of images.
is_training: `bool` if True, the model is in training mode.
Returns:
a `dict` containing `int` keys for continuous feature levels [2, 3, 4, 5].
The values are corresponding feature hierarchy in ResNet with shape
[batch_size, height_l, width_l, num_filters].
"""
with backend.get_graph().as_default():
with tf.name_scope('resnet%s' % self._resnet_depth):
return self._resnet_fn(inputs, is_training)
def __call__(self, roi_features, is_training=None):
"""Box and class branches for the Mask-RCNN model.
Args:
roi_features: A ROI feature tensor of shape
[batch_size, num_rois, height_l, width_l, num_filters].
is_training: `boolean`, if True if model is in training mode.
Returns:
class_outputs: a tensor with a shape of
[batch_size, num_rois, num_classes], representing the class predictions.
box_outputs: a tensor with a shape of
[batch_size, num_rois, num_classes * 4], representing the box
predictions.
"""
with backend.get_graph().as_default(), tf.name_scope('fast_rcnn_head'):
# reshape inputs beofre FC.
_, num_rois, height, width, filters = roi_features.get_shape(
).as_list()
roi_features = tf.reshape(roi_features,
[-1, num_rois, height * width * filters])
net = tf.keras.layers.Dense(units=self._mlp_head_dim,
activation=None,
name='fc6')(roi_features)
net = self._batch_norm_relu(fused=False)(net,
is_training=is_training)
net = tf.keras.layers.Dense(units=self._mlp_head_dim,
activation=None,
name='fc7')(net)
net = self._batch_norm_relu(fused=False)(net,
is_training=is_training)
class_outputs = tf.keras.layers.Dense(
self._num_classes,
kernel_initializer=tf.keras.initializers.RandomNormal(
stddev=0.01),
bias_initializer=tf.zeros_initializer(),
name='class-predict')(net)
box_outputs = tf.keras.layers.Dense(
self._num_classes * 4,
kernel_initializer=tf.keras.initializers.RandomNormal(
stddev=0.001),
bias_initializer=tf.zeros_initializer(),
name='box-predict')(net)
return class_outputs, box_outputs
def call_context(layer):
"""Scope that marks when we are currently inside a Layer/Model's `call`."""
was_in_call = is_in_call_context()
was_frozen = is_in_frozen_context()
was_in_keras_graph = getattr(_call_context, 'in_keras_graph', False)
_call_context.in_call = True
_call_context.in_keras_graph = (
was_in_keras_graph or
getattr(backend.get_graph(), 'name', None) == 'keras_graph')
if not layer.trainable:
_call_context.frozen = True
try:
yield
finally:
_call_context.in_call = was_in_call
_call_context.frozen = was_frozen
_call_context.in_keras_graph = was_in_keras_graph
def get_gradients_for_keras(optimizer, loss, params):
from tensorflow.python.util import nest
from tensorflow.python.keras import backend
from tensorflow.python.ops import gradients
from tensorflow.python.ops import clip_ops
from tensorflow.python.keras.optimizers import TFOptimizer
params = nest.flatten(params)
if isinstance(optimizer, TFOptimizer):
scope_name = optimizer.optimizer._name
else:
scope_name = optimizer._name
with backend.get_graph().as_default(), backend.name_scope(scope_name +
"/gradients"):
grads = gradients.gradients(loss, params)
all_reduced_grads = []
for grad, param in zip(grads, params):
if grad is None:
raise ValueError(
"Variable {} has `None` for gradient. "
"Please make sure that all of your ops have a "
"gradient defined (i.e. are differentiable). "
"Common ops without gradient: "
"K.argmax, K.round, K.eval.".format(param))
grad = process_grad(grad)
with tf.control_dependencies([param]):
grad_i = tf.identity(grad, name="zoo_identity_op_for_grad")
all_reduced_grads.append(grad_i)
grads = all_reduced_grads
if hasattr(optimizer, "clipnorm"):
grads = [
clip_ops.clip_by_norm(g, optimizer.clipnorm) for g in grads
]
if hasattr(optimizer, "clipvalue"):
grads = [
clip_ops.clip_by_value(g, -optimizer.clipvalue,
optimizer.clipvalue) for g in grads
]
return grads
def __call__(self, inputs, is_training=None):
with backend.get_graph().as_default():
model = SpineNet(
min_level=self._min_level,
max_level=self._max_level,
block_specs=self._block_specs,
endpoints_num_filters=self._endpoints_num_filters,
resample_alpha=self._resample_alpha,
block_repeats=self._block_repeats,
filter_size_scale=self._filter_size_scale,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activation=self._activation,
use_sync_bn=self._use_sync_bn,
norm_momentum=self._norm_momentum,
norm_epsilon=self._norm_epsilon)
return model(inputs)
def fit(self, train_set, test_set):
logdir = os.path.join(self.params.logdir, self.model.name)
if not os.path.exists(logdir):
os.makedirs(logdir)
train_writer = tf.summary.create_file_writer(
os.path.join(logdir, 'train'))
test_writer = tf.summary.create_file_writer(
os.path.join(logdir, 'test'))
with train_writer.as_default():
summary_ops_v2.graph(K.get_graph(), step=0)
do_callbacks('on_train_begin', self.model.callbacks)
for epoch in range(self.init_epoch, self.params.training.epochs):
self.train_one_epoch(train_set, train_writer, epoch)
self.evaluate(test_set, test_writer, epoch)
self.checkpoint(self.manager, epoch + 1,
self.params.training.save_frequency)
self.checkpoint(self.manager, self.params.training.epochs, 0)
def get_gradients_hessian(self, loss, params):
"""Returns gradients and Hessian of `loss` with respect to `params`.
Arguments:
loss: Loss tensor.
params: List of variables.
Returns:
List of gradient and Hessian tensors.
Raises:
ValueError: In case any gradient cannot be computed (e.g. if gradient
function not implemented).
"""
params = nest.flatten(params)
with backend.get_graph().as_default(), backend.name_scope(
self._name + "/gradients"):
grads = gradients.gradients(loss, params)
for grad, param in zip(grads, params):
if grad is None:
raise ValueError(
"Variable {} has `None` for gradient. "
"Please make sure that all of your ops have a "
"gradient defined (i.e. are differentiable). "
"Common ops without gradient: "
"K.argmax, K.round, K.eval.".format(param))
# WARNING: for now we do not support gradient clip
# grads = self._clip_gradients(grads)
v = [np.random.uniform(0, 1, size=p.shape) for p in params]
for vi in v:
vi[vi < 0.5] = -1
vi[vi >= 0.5] = 1
v = [tf.convert_to_tensor(vi, dtype=tf.dtypes.float32) for vi in v]
vprod = tf.reduce_sum(
[tf.reduce_sum(vi * grad) for vi, grad in zip(v, grads)])
Hv = gradients.gradients(vprod, params)
Hd = [tf.abs(Hvi * vi) for Hvi, vi in zip(Hv, v)]
return grads, Hd
def set_model(self, model):
"""Sets Keras model and writes graph if specified."""
self.model = model.model
with context.eager_mode():
self._close_writers()
if self.write_graph:
with self._get_writer(self._train_run_name).as_default():
with summary_ops_v2.always_record_summaries():
if not self.model.run_eagerly:
summary_ops_v2.graph(K.get_graph(), step=0)
summary_writable = (
self.model._is_graph_network or # pylint: disable=protected-access
self.model.__class__.__name__ == 'Sequential') # pylint: disable=protected-access
if summary_writable:
summary_ops_v2.keras_model('keras',
self.model,
step=0)
if self.embeddings_freq:
self._configure_embeddings()
def _wait_for_variable_initialization(session):
"""Utility to wait for variables to be initialized."""
all_variables = K._get_variables(K.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(
[variables.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
def _wait_for_variable_initialization(session):
"""Utility to wait for variables to be initialized."""
all_variables = K._get_variables(K.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(
[variables.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
def get_trainable_variables():
"""Get a list of trainable TensorFlow variables.
Parameters
----------
train_only : boolean
If True, only get the trainable variables.
Returns
-------
list of Tensor
A list of trainable TensorFlow variables
Examples
--------
"""
if KERAS_MODELS or LooseVersion(tf.__version__) >= LooseVersion("2.0.0"):
logging.warning(
"In TF2.x, only trainable variables created with Keras Models are captured for logging.\n"
"In TF1.x, if any keras model is defined. Only variables created inside Keras Models will be logged."
)
var_list = list()
for model in KERAS_MODELS:
var_list.extend(model.trainable_variables)
# Keep only a list of unique variables (remove potential duplicates)
var_list = list(set(var_list))
# clearing the list of Keras Model to avoid memory leaks
KERAS_MODELS.clear()
return [var for var in sorted(var_list, key=lambda v: v.name)]
else:
# return tf.trainable_variables() # deprecated in TF2.x
from tensorflow.python.keras.backend import get_graph
return get_graph().get_collection('trainable_variables')
def _make_eager_execution_function(model, mode):
"""Makes function to run one step of distributed model eager execution."""
def _per_device_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 = K.get_graph()
with global_graph.as_default(), strategy.scope():
# First we gather the relevant portions of the model across all replicas.
# `K._scratch_graph(global_graph)` signals to Keras that it should not
# lift to a separate graph when creating the per-replica functions.
with K._scratch_graph(global_graph):
# Create train ops on each of the devices when we call
# `_per_device_fit_function`.
grouped = strategy.extended.call_for_each_replica(
_per_device_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 K.function(
all_inputs,
all_outputs,
name='eager_distributed_{}_function'.format(mode))
def __init__(self,
input_shape=None,
batch_size=None,
dtype=None,
input_tensor=None,
sparse=False,
name=None,
**kwargs):
strategy = distribution_strategy_context.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 value {} cannot be '
'divisible by 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 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 = sparse
self.batch_size = batch_size
self.supports_masking = True
if isinstance(input_shape, tensor_shape.TensorShape):
input_shape = tuple(input_shape.as_list())
elif isinstance(input_shape, int):
input_shape = (input_shape,)
if 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():
# In graph mode, create a graph placeholder to call the layer on.
if sparse:
input_tensor = backend.placeholder(
shape=batch_input_shape,
dtype=dtype,
name=self.name,
sparse=True)
else:
input_tensor = backend.placeholder(
shape=batch_input_shape,
dtype=dtype,
name=self.name)
self.is_placeholder = True
self._batch_input_shape = batch_input_shape
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
self._batch_input_shape = tuple(input_tensor.shape.as_list())
# Create an input node to add to self.outbound_node
# and set output_tensors' _keras_history.
input_tensor._keras_history = (self, 0, 0) # pylint: disable=protected-access
input_tensor._keras_mask = None
base_layer.Node(
self,
inbound_layers=[],
node_indices=[],
tensor_indices=[],
input_tensors=[input_tensor],
output_tensors=[input_tensor])
def __init__(self,
input_shape=None,
batch_size=None,
dtype=None,
input_tensor=None,
sparse=False,
name=None,
**kwargs):
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 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)
super(InputLayer, self).__init__(dtype=dtype, name=name)
self.built = True
self.sparse = sparse
self.batch_size = batch_size
self.supports_masking = True
if isinstance(input_shape, tensor_shape.TensorShape):
input_shape = tuple(input_shape.as_list())
if 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():
# In graph mode, create a graph placeholder to call the layer on.
if sparse:
input_tensor = backend.placeholder(
shape=batch_input_shape,
dtype=dtype,
name=self.name,
sparse=True)
else:
input_tensor = backend.placeholder(
shape=batch_input_shape,
dtype=dtype,
name=self.name)
self.is_placeholder = True
self._batch_input_shape = batch_input_shape
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
self._batch_input_shape = tuple(input_tensor.get_shape().as_list())
# Create an input node to add to self.outbound_node
# and set output_tensors' _keras_history.
input_tensor._keras_history = (self, 0, 0) # pylint: disable=protected-access
base_layer.Node(
self,
inbound_layers=[],
node_indices=[],
tensor_indices=[],
input_tensors=[input_tensor],
output_tensors=[input_tensor])