def complete_graph(self,
graph: tf_compat.Graph = None,
sess: tf_compat.Session = None):
"""
Complete modifying the graph. Should be called after modifying is complete.
Cleans up any ops that should be removed or reordered.
: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)
if not graph:
graph = tf_compat.get_default_graph()
if not sess:
sess = tf_compat.get_default_session()
for mod in self.modifiers:
mod.complete_graph(graph, sess)
return graph
def resnet_v2_50(init_weights):
tf_compat.reset_default_graph()
image_size = 224
inputs = tf_compat.random_normal([1, image_size, image_size, 3])
with slim.arg_scope(resnet_v2.resnet_arg_scope()):
logits, _ = resnet_v2.resnet_v2_50(inputs, 1000, is_training=False)
return tf_compat.get_default_graph()
def simple_matmul_net(init_weights):
tf_compat.reset_default_graph()
n_inputs = 28 * 28
n_hidden1 = 300
n_hidden2 = 100
n_outputs = 10
X = tf_compat.placeholder(tf_compat.float32, shape=(None, n_inputs), name="X")
def neuron_layer(X, n_neurons, name, activation=None):
with tf_compat.name_scope(name):
n_inputs = int(X.get_shape()[1])
stddev = 2 / np.sqrt(n_inputs)
init = tf_compat.truncated_normal((n_inputs, n_neurons), stddev=stddev)
W = tf_compat.Variable(init, name="kernel")
b = tf_compat.Variable(tf_compat.zeros([n_neurons]), name="bias")
Z = tf_compat.matmul(X, W) + b
if activation is not None:
return activation(Z)
else:
return Z
with tf_compat.name_scope("dnn"):
hidden1 = neuron_layer(
X, n_hidden1, name="hidden1", activation=tf_compat.nn.relu
)
hidden2 = neuron_layer(
hidden1, n_hidden2, name="hidden2", activation=tf_compat.nn.relu
)
neuron_layer(hidden2, n_outputs, name="outputs")
return tf_compat.get_default_graph()
def test_modify_estimator(
self,
modifier_lambda: Callable[[], Modifier],
graph_lambda: Callable[[], tf_compat.Graph],
steps_per_epoch: int,
):
def model_fn(features, labels, mode, params):
graph_lambda()
return tf_compat.estimator.EstimatorSpec(mode)
modifier = modifier_lambda()
tf_compat.get_default_graph()
estimator = tf_compat.estimator.Estimator(model_fn=model_fn, )
assert estimator._model_fn == model_fn
modifier.modify_estimator(estimator, steps_per_epoch)
assert estimator._model_fn != model_fn
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 simple_conv2d_net(init_weights):
tf_compat.reset_default_graph()
X = tf_compat.placeholder(tf_compat.float32, [None, 32, 40, 1])
W = tf_compat.Variable(
tf_compat.convert_to_tensor(init_weights, dtype=tf_compat.float32)
)
b = tf_compat.Variable(tf_compat.random_normal([64]), dtype=tf_compat.float32)
conv1 = tf_compat.nn.conv2d(X, W, strides=[1, 1, 1, 1], padding="VALID")
conv1 = tf_compat.nn.bias_add(conv1, b)
conv1 = tf_compat.nn.max_pool(
conv1, ksize=[1, 1, 3, 1], strides=[1, 1, 1, 1], padding="VALID"
)
return tf_compat.get_default_graph()
def _model_func(
features: Dict[str, tf_compat.Tensor],
labels: Dict[str, tf_compat.Tensor],
mode: tf_compat.estimator.ModeKeys,
params: Dict[str, Any],
):
spec = orig_model_func(features=features,
labels=labels,
mode=mode,
params=params)
graph = tf_compat.get_default_graph()
with graph.as_default():
global_step = tf_compat.train.get_or_create_global_step()
mod_ops, mod_extras = self.create_ops(steps_per_epoch, global_step,
graph)
hook = ModifierSessionRunHook(self, steps_per_epoch, mod_ops,
mod_extras)
replace_kwargs = {}
if mode == tf_compat.estimator.ModeKeys.TRAIN:
replace_kwargs = {"training_hooks": [hook]}
if spec.training_hooks:
replace_kwargs["training_hooks"].extend(
spec.training_hooks)
orig_saver = spec.scaffold.saver
saver = tf_compat.train.Saver(
var_list=None,
reshape=orig_saver._reshape,
sharded=orig_saver._sharded,
max_to_keep=orig_saver._max_to_keep,
keep_checkpoint_every_n_hours=orig_saver.
_keep_checkpoint_every_n_hours,
name=orig_saver._name,
restore_sequentially=orig_saver._restore_sequentially,
pad_step_number=orig_saver._pad_step_number,
save_relative_paths=orig_saver._save_relative_paths,
filename=orig_saver._filename,
)
replace_kwargs["scaffold"] = tf_compat.train.Scaffold(
saver=saver, copy_from_scaffold=spec.scaffold)
spec = spec._replace(**replace_kwargs)
return spec
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