def create_ks_scheduled_constant_graph_ops( graph: tf_compat.Graph, global_step: tf_compat.Variable, var_names: List[str], begin_step: int, end_step: int, ks_group: str, ) -> Tuple[tf_compat.Tensor, List[PruningOpVars]]: """ Creates constant model pruning ops. Does not modify the graph. :param graph: the tf graph to pull the operator out of for applying the pruning to :param global_step: the global optimizer step for the training graph :param var_names: a list of names or regex patterns to create constant ops for within the graph :param begin_step: the global step to begin pruning at :param end_step: the global step to end pruning at :param ks_group: the group identifier the scope should be created under :return: a tuple containing the update operation to run in a session, a list of the pruning ops and vars for each desired op in the graph """ pruning_op_vars = [] is_start_step = tf_compat.equal(global_step, begin_step) is_end_step = tf_compat.equal(global_step, end_step) for op, op_input in get_ops_and_inputs_by_name_or_regex(var_names, graph): op_vars = create_constant_op_pruning(op, op_input, is_start_step, is_end_step, ks_group) pruning_op_vars.append(op_vars) update_op = get_scheduled_update_op(pruning_op_vars, ks_group) return update_op, pruning_op_vars
def create_graph_ops_pruning( graph: tf_compat.Graph, var_names: List[str], sparsity: tf_compat.Tensor, update_ready: tf_compat.Tensor, leave_enabled: bool, is_after_end_step: tf_compat.Tensor, ks_group: str, mask_creator: PruningMaskCreator, ) -> List[PruningOpVars]: """ Creates the necessary variables and operators to gradually apply sparsity to a given list of operators in a graph. Handles setting a mask on an operator to the given sparsity. Sets the mask based on pruning away the lowest absolute magnitude weights. :param graph: the tf graph to pull the operator out of for applying the pruning to :param var_names: the names or regex patterns of names of variables to prune in the graph to the given sparsity :param sparsity: the target sparsity to use for assigning the masks :param update_ready: the tensor where if true will update the mask from sparsity, if false will not update the mask :param leave_enabled: True to continue masking the weights after end_epoch, False to stop masking :param is_after_end_step: tensor that is true if the current global step is after end_epoch :param ks_group: the group identifier the scope should be created under :param mask_creator: optional object to define sparisty mask creation :return: a list of the created named tuples each containing the assignment op, mask variable, threshold tensor, and masked tensor """ pruning_op_vars = [] variable_masks = {} # cache of mask vars for input variables for op, op_input in get_ops_and_inputs_by_name_or_regex(var_names, graph): if op_input not in variable_masks: op_vars = create_op_pruning( op, op_input, sparsity, update_ready, leave_enabled, is_after_end_step, ks_group, mask_creator, ) pruning_op_vars.append(op_vars) variable_masks[op_input] = op_vars else: # Reuse masks if the input variable is shared and already computed _, _, mask_update, mask, masked = variable_masks[op_input] pruning_op_vars.append( PruningOpVars(op, op_input, mask_update, mask, masked)) tf_compat.add_to_collection( PruningScope.collection_name(ks_group, PruningScope.OPS), op) tf_compat.add_to_collection( PruningScope.collection_name(ks_group, PruningScope.OPS_INPUT), op_input) return pruning_op_vars
def pruning_loss_sens_magnitude( graph: tf_compat.Graph = None, sess: tf_compat.Session = None, sparsity_levels: Union[List[float], Tuple[float, ...]] = default_pruning_sparsities_loss(True), ) -> PruningLossSensitivityAnalysis: """ Approximated kernel sparsity (pruning) loss analysis for a given model. Returns the results for each prunable param (conv, linear) in the model. Approximated by taking the magnitudes of the weights. :param graph: the graph to inject pruning ops and vars into, if not supplied uses get_default_graph() :param sess: the session to use :param sparsity_levels: the sparsity levels to calculate the loss for for each param :return: the analysis results for the model """ if not graph: graph = tf_compat.get_default_graph() if not sess: sess = tf_compat.get_default_session() prunable_ops_and_inputs = get_ops_and_inputs_by_name_or_regex(["re:.*"], graph) analysis = PruningLossSensitivityAnalysis() for op_index, (_, op_tens) in enumerate(prunable_ops_and_inputs): weight = sess.run(op_tens) values = numpy.sort(numpy.abs(weight.reshape(-1))) prev_index = 0 for sparsity in sparsity_levels: val_index = round(sparsity * len(values)) if val_index >= len(values): val_index = len(values) - 1 if sparsity <= 1e-9: baseline = True sparsity = 0.0 sparse_avg = 0.0 else: baseline = False if val_index > prev_index: sparse_avg = values[prev_index:val_index].mean().item() prev_index = val_index else: sparse_avg = values[val_index].item() prev_index = val_index + 1 analysis.add_result(None, op_tens.name, op_index, sparsity, sparse_avg, baseline) return analysis
def pruning_loss_sens_op_vars( graph: tf_compat.Graph = None, var_names: Union[List[str], Tuple[str]] = ("re:.*", ), mask_type: Union[str, List[int], PruningMaskCreator] = "unstructured", ) -> List[SparsePruningOpVars]: """ Edit the graph for to inject pruning ops and vars to allow for a ks loss sensitivity analysis. Note: this must be run outside of a session for it to take effect. :param graph: the graph to inject pruning ops and vars into, if not supplied uses get_default_graph() :param var_names: List of variable names or regex patterns of variables to get the op vars for. Defaults to matching all variables :param mask_type: String to define type of sparsity (options: ['unstructured', 'channel', 'filter']), List to define block shape of a parameter's in and out channels, or a SparsityMaskCreator object. default is 'unstructured' :return: the created pruning op vars to be used in approx_ks_loss_sensitivity and one_shot_ks_loss_sensitivity """ if not graph: graph = tf_compat.get_default_graph() mask_creator = mask_type if not isinstance(mask_type, PruningMaskCreator): mask_creator = load_mask_creator(mask_type) ks_group = pruning_loss_sens_one_shot.__name__ prunable_ops_and_inputs = get_ops_and_inputs_by_name_or_regex( var_names, graph) op_vars = [] with graph.as_default(): for prune_op, prune_op_input in prunable_ops_and_inputs: with tf_compat.name_scope( PruningScope.model(prune_op, ks_group, trailing_slash=True)): sparsity = tf_compat.placeholder(dtype=tf_compat.float32, name="sparsity_placeholder") update = tf_compat.constant(True, tf_compat.bool) prune_op_var = create_op_pruning( prune_op, prune_op_input, sparsity, update, True, None, ks_group, mask_creator, ) op_vars.append(SparsePruningOpVars(prune_op_var, sparsity)) return op_vars
def test_get_ops_and_inputs_by_name_or_regex( net_const, var_names, expected_ops, expected_tens, ): with tf_compat.Graph().as_default() as graph: net_const() ops_and_inputs = get_ops_and_inputs_by_name_or_regex(var_names, graph) assert len(ops_and_inputs) == len(expected_ops) for op, inp in ops_and_inputs: assert op.name in expected_ops assert clean_tensor_name(inp.name) in expected_tens
def create_ops( self, steps_per_epoch: int, global_step: tf_compat.Tensor, graph: tf_compat.Graph, ) -> Tuple[List[Union[tf_compat.Tensor, tf_compat.Operation]], Dict[str, Any]]: """ Create the sparsity ops to modify the training graph according to the settings for the current instance. :param steps_per_epoch: the number of steps (batches) per training epoch :param global_step: the global step used while training :param graph: the graph to be modified :return: a tuple (list of ops, dict of named ops / tensors) to be run or used for modifying the training process. """ mod_ops, mod_extras = super().create_ops(graph, None, None) start_step, end_step = self.start_end_steps(steps_per_epoch, after_optim=True) params = ( self._params if self._params != ALL_TOKEN else [ clean_tensor_name(var.name) for _, var in # Have ALL_TOKEN match to all variable names for now get_ops_and_inputs_by_name_or_regex(["re:.*"], graph) ] ) with graph.as_default(): update_op, prune_op_vars = create_ks_scheduled_constant_graph_ops( graph, global_step, params, start_step, end_step, self.ks_group, ) if self.log_types == ALL_TOKEN or "tensorboard" in self.log_types: mod_extras[EXTRAS_KEY_SUMMARIES] = create_summaries_pruning( prune_op_vars ) mod_ops.append(update_op) self._prune_op_vars = prune_op_vars # self._update_ready = tf_compat.constant(False, name="nm_update_ready") return mod_ops, mod_extras
def create_ops( self, steps_per_epoch: int, global_step: tf_compat.Tensor, graph: tf_compat.Graph, ) -> Tuple[List[Union[tf_compat.Tensor, tf_compat.Operation]], Dict[str, Any]]: """ Create the sparsity ops to modify the training graph according to the settings for the current instance. :param steps_per_epoch: the number of steps (batches) per training epoch :param global_step: the global step used while training :param graph: the graph to be modified :return: a tuple (list of ops, dict of named ops / tensors) to be run or used for modifying the training process. """ mod_ops, mod_extras = super().create_ops(graph, steps_per_epoch, global_step) start_step, end_step = self.start_end_steps(steps_per_epoch, after_optim=True) update_frequency_step = self.update_frequency_steps(steps_per_epoch) params = ( self._params if self._params != ALL_TOKEN else [ clean_tensor_name(var.name) for _, var in # Have ALL_TOKEN match to all variable names for now get_ops_and_inputs_by_name_or_regex(["re:.*"], graph) ]) with graph.as_default(): ( update_op, prune_op_vars, update_ready, sparsity, ) = get_or_create_ks_scheduled_graph_ops( graph, global_step, params, start_step, end_step, update_frequency_step, self._init_sparsity, self._final_sparsity, self.exponent, self._leave_enabled, self.ks_group, self._mask_creator, ) if self.log_types == ALL_TOKEN or "tensorboard" in self.log_types: mod_extras[EXTRAS_KEY_SUMMARIES] = create_summaries_pruning( prune_op_vars) mod_ops.append(update_op) self._prune_op_vars = prune_op_vars self._update_ready = update_ready self._sparsity = sparsity # Create and cache the mask initializers to be run # through initialize_session. When using the estimator, # the initialization is done as part of the init_fn of # the training scaffold object, at which the graph cannot # be changed (hence the creation and caching) masks = [op_vars.mask for op_vars in self._prune_op_vars] self._mask_initializer = (tf_compat.variables_initializer(masks) if masks else None) return mod_ops, mod_extras