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