def _spec_from_ratio(self, net_sens, ratio, excludes=[]):
inputs = utils.dummy_inputs(self._input_specs)
flops, _ = summary.model_complexity(self._module, inputs)
expected_flops = (1 - ratio) * flops
flops_tolerance = 1e-2
min_th = 1e-5
max_th = 1 - min_th
num_attempts = 0
max_attempts = 100
prev_spec = None
cur_spec = None
while num_attempts < max_attempts:
prev_spec = cur_spec
num_attempts += 1
threshold = (min_th + max_th) / 2
cur_spec = self._spec_from_threshold(net_sens, threshold, excludes)
if prev_spec and prev_spec == cur_spec:
continue
pruned_model, _, _ = self._prune(self._graph, cur_spec)
current_flops, _ = summary.model_complexity(pruned_model, inputs)
error = abs(expected_flops - current_flops) / expected_flops
if error < flops_tolerance:
break
if current_flops < expected_flops:
max_th = threshold
else:
min_th = threshold
logging.info(
'original_flops: {}, expected_flops: {}, current_flops: {}'.format(
flops, expected_flops, current_flops))
return cur_spec
def _ana(self, eval_fn, args=()):
analyser = ana_lib.ModelAnalyser(self._graph)
steps = analyser.steps()
for step in range(steps):
spec = analyser.spec(step)
model, _, _ = self._prune(self._graph, spec)
model.eval()
model = model.cuda()
score = eval_fn(model, *args).item()
analyser.record(step, score)
logging.info('Analysis complete %d/%d' % (cur_step + 1, steps))
analyser.save()
def prune(self,
ratio=None,
threshold=None,
excludes=[],
output_script='graph.py'):
"""Prune the network by given ratio or threshold.
Arguments:
ratio: The expected percentage of FLOPs reduction. This is just a hint
value, the actual FLOPs drop not necessarily strictly to this value
after pruning.
threshold: Relative proportion of model performance loss
that can be tolerated.
excludes: Modules that need to prevent from pruning.
output_script: Filepath that saves the generated script
used for rebuilding model.
Return:
A `PruningModule` object works like a normal torch.nn.Module with
addtional pruning info.
"""
sens_path = pruning_lib.sens_path(self._graph)
if not os.path.exists(sens_path):
raise RuntimeError("Must call ana() before runnig prune.")
net_sens = pruning_lib.read_sens(sens_path)
excluded_nodes = []
if excludes:
module_to_node = utils.map_original_module_to_node(
self._module, self._graph)
for module in excludes:
excluded_nodes.append(module_to_node[id(module)])
if threshold:
logging.info('start pruning by threshold = {}'.format(threshold))
spec = self._spec_from_threshold(net_sens, threshold,
excluded_nodes)
elif ratio:
logging.info('start pruning by ratio = {}'.format(ratio))
spec = self._spec_from_ratio(net_sens, ratio, excluded_nodes)
else:
raise ValueError("One of 'ratio' or 'threshold' must be given.")
pruned_model, pruned_graph, pruning_info = self._prune(
self._graph, spec, output_script)
return PruningModule(pruned_model, pruned_graph, pruning_info)
def __call__(self, input_queue, output_queue):
# We have to reparse graph here to recreate the global variable
# _NNDCT_OP_2_TORCH_OP. (obviously not a good idea)
# TorchScriptWriter use this global map to generate the script.
graph = parse_to_graph(self.module, self.input_specs)
pruner = pruner_lib.ChannelPruner(graph)
analyser = ana_lib.ModelAnalyser(graph)
steps = analyser.steps()
while not input_queue.empty():
cur_step = input_queue.get()
spec = analyser.spec(cur_step)
pruned_graph, _ = pruner.prune(spec)
rebuilt_module, _ = utils.rebuild_module(pruned_graph)
module = rebuilt_module.cuda()
module.eval()
score = self.eval_fn(module, *self.args).item()
output_queue.put((cur_step, score))
logging.info('Analysis complete %d/%d' % (cur_step + 1, steps))