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 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])