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 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 complete_graph(self, graph: tf_compat.Graph, sess: tf_compat.Session):
"""
Complete modifying the graph.
Resets the pruned op's variables using the created masks to zero out
the pruned weights for saving.
:param graph: the modified graph that should be completed and cleaned.
if not supplied, then will use the default graph
:param sess: the session to use for completing the modified graph.
if not supplied, then will use the default session
:return: the cleaned graph
"""
super().complete_graph(graph, sess)
with graph.as_default():
apply_op_vars_masks(self.prune_op_vars, self.ks_group, sess)
def create_ops(
self,
steps_per_epoch: int,
global_step: tf_compat.Variable,
graph: tf_compat.Graph,
) -> Tuple[List[Union[tf_compat.Tensor, tf_compat.Operation]], Dict[str, Any]]:
"""
Create switch case computing the learning rate at a given global step and
extras created by individual LR modifiers
:param steps_per_epoch: the number of steps per training epoch
:param global_step: the global step used while training
:param graph: the graph to be modified
:return: a tuple (list of empty ops, dict of named ops/tensors for learning
rate and summaries as extras)
"""
mod_ops, mod_extras = super().create_ops(graph, steps_per_epoch, global_step)
name_scope = "{}/{}".format(NM_RECAL, self.__class__.__name__)
with graph.as_default():
with tf_compat.name_scope(name_scope):
pred_fn_pairs = []
global_step = tf_compat.cast(global_step, tf_compat.int64)
for index, child in enumerate(self._lr_modifiers):
with tf_compat.name_scope(str(index)):
_, child_extras = child.create_ops(
steps_per_epoch, global_step, graph
)
child_lr = child_extras[EXTRAS_KEY_LEARNING_RATE]
child_start_step, _ = child.start_end_steps(
steps_per_epoch, after_optim=False
)
child_select = tf_compat.greater_equal(
global_step,
tf_compat.constant(child_start_step, tf_compat.int64),
name="active",
)
pred_fn_pairs.append((child_select, lambda lr=child_lr: lr))
learning_rate = tf_compat.case(pred_fn_pairs)
_add_lr_extras(mod_extras, learning_rate, self.log_types)
return mod_ops, mod_extras
def create_ops(
self,
steps_per_epoch: int,
global_step: tf_compat.Tensor = None,
graph: tf_compat.Graph = None,
) -> Tuple[List[Union[tf_compat.Tensor, tf_compat.Operation]], Dict[str,
Any]]:
"""
Create modifying operations and tensors in the graph.
| Returns a tuple containing:
| - modifying ops that should be run in a session on each global step.
| - named extras (ops / tensors) created in the graph that can be used
| by other ops such as a learning rate for the optimizer
:param steps_per_epoch: the number of steps (batches) per training epoch
:param global_step: the global step used while training.
if not supplied, then will use get_or_create_global_step()
:param graph: the graph to be modified,
if not supplied, then will use the default graph
:return: a tuple (list of ops, dict of named ops / tensors)
to be run or used for modifying the training process
"""
if not graph:
graph = tf_compat.get_default_graph()
if not global_step:
with graph.as_default():
global_step = tf_compat.train.get_or_create_global_step()
ops, extras = super().create_ops(steps_per_epoch, global_step, graph)
mod_ops_extras = [
mod.create_ops(steps_per_epoch, global_step, graph)
for mod in self.modifiers
] # List[Tuple[List, Dict]]
merged_ops = list(
itertools.chain.from_iterable(
[_ops for (_ops, _) in mod_ops_extras if _ops]))
mod_extras = [_extras for (_, _extras) in mod_ops_extras
] # List[Dict[str, Any]]
extras = {}
for _extras in mod_extras:
for key, val in _extras.items():
if key not in extras:
extras[key] = val
else:
# This key exists before either as a list or a single value
if not isinstance(extras[key], List):
raise ValueError(
"extras[{}] has been recorded with unique "
"value and cannot be merged".format(key))
if not isinstance(val, List):
raise ValueError(
"extras[{}] has been recorded as list, "
"requiring new list to merge".format(key))
extras[key].extend(val)
with tf_compat.name_scope(NM_RECAL):
ops.append(
tf_compat.group(merged_ops,
name=ScheduledModifierManager.RECAL_UPDATE))
return ops, 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