def get_normalized_recipe(recipe):
return ThinningRecipe(
modules={
utils.normalize_module_name(k): v
for k, v in recipe.modules.items()
},
parameters={
utils.normalize_module_name(k): v
for k, v in recipe.parameters.items()
},
)
def successors_f(self,
node_name,
successors_types,
done_list=None,
logging=None,
denorm_names=True):
"""Returns a list of <op>'s successors, if they match the <successors_types> criteria.
Traverse the graph, starting at node <node_name>, and search for successor
nodes, that have one of the node types listed in <successors_types>.
If none is found, then return an empty list.
<node_name> and the returned list of successors are strings, because
"""
node_name = utils.normalize_module_name(node_name)
node = self.find_op(node_name)
node_is_an_op = True
if node is None:
node_is_an_op = False
node = self.find_param(node_name)
if node is None:
msglogger.warning(
"successors_f: Could not find node {}".format(node_name))
return []
if done_list is None:
done_list = []
done_list.append(node_name)
if not isinstance(successors_types, list):
successors_types = [successors_types]
if node_is_an_op:
# We check if we found the type of node we're looking for,
# and that this is not the first node in our search.
if node['type'] in successors_types and len(done_list) > 1:
return [
utils.denormalize_module_name(self._src_model, node_name)
if denorm_names else node_name
]
# This is an operation node
succs = [
edge.dst for edge in self.edges
if (edge.src == node_name and edge.dst not in done_list)
]
else:
# This is a data node
succs = [
edge.dst for edge in self.edges
if (edge.src == node_name and edge.dst not in done_list)
]
ret = []
for successor in succs:
ret += self.successors_f(successor, successors_types, done_list,
logging, denorm_names)
return ret
def named_params_layers(self):
for param_name, param in self._src_model.named_parameters():
# remove the extension of param_name, and then normalize it
# to create a normalized layer name
normalized_layer_name = utils.normalize_module_name('.'.join(
param_name.split('.')[:-1]))
sgraph_layer_name = utils.denormalize_module_name(
self._src_model, normalized_layer_name)
yield sgraph_layer_name, param_name, param
def predecessors_f(self,
node_name,
predecessors_types,
done_list=None,
logging=None,
denorm_names=True):
"""Returns a list of <op>'s predecessors, if they match the <predecessors_types> criteria.
"""
node_name = utils.normalize_module_name(node_name)
node = self.find_op(node_name)
node_is_an_op = True
if node is None:
node_is_an_op = False
node = self.find_param(node_name)
if node is None:
msglogger.warning(
"predecessors_f: Could not find node {}".format(node_name))
return []
if done_list is None:
done_list = []
done_list.append(node_name)
if not isinstance(predecessors_types, list):
predecessors_types = [predecessors_types]
if node_is_an_op:
# We check if we found the type of node we're looking for,
# and that this is not the first node in our search.
if node['type'] in predecessors_types and len(done_list) > 1:
return [
utils.denormalize_module_name(self._src_model, node_name)
if denorm_names else node_name
]
# This is an operation node
preds = [
edge.src for edge in self.edges
if (edge.dst == node_name and edge.src not in done_list)
]
else:
# This is a data node
preds = [
edge.src for edge in self.edges
if (edge.dst == node_name and edge.src not in done_list)
]
ret = []
for predecessor in preds:
ret += self.predecessors_f(predecessor, predecessors_types,
done_list, logging, denorm_names)
return ret
def get_model_compute_budget(model, dataset, layers_to_prune=None):
"""Return the compute budget of the Convolution layers in an image-classifier.
"""
dummy_input = utils.get_dummy_input(dataset)
g = SummaryGraph(model, dummy_input)
total_macs = 0
for name, m in model.named_modules():
if isinstance(m, torch.nn.Conv2d):
# Use the SummaryGraph to obtain some other details of the models
conv_op = g.find_op(normalize_module_name(name))
assert conv_op is not None
total_macs += conv_op['attrs']['MACs']
del g
return total_macs
def load_state_dict(self, state, normalize_dataparallel_keys=False):
try:
loaded_masks = state['masks_dict']
except KeyError as exception:
msglogger.error('could not load the CompressionScheduler state.'
' masks_dict is missing from state')
with contextlib.suppress(TypeError):
msglogger.debug('Scheduler state keys are: {}'.format(', '.join(state)))
raise
if normalize_dataparallel_keys:
loaded_masks = {normalize_module_name(k): v for k, v in loaded_masks.items()}
device = model_device(self.model)
for name, mask in self.zeros_mask_dict.items():
masker = self.zeros_mask_dict[name]
masker.mask = loaded_masks[name]
if masker.mask is not None:
masker.mask = masker.mask.to(device)
def log_model_buffers(self, model, buffer_names, tag_prefix, epoch,
completed, total, freq):
"""Logs values of model buffers.
Notes:
1. Each buffer provided is logged in a separate CSV file
2. Each CSV file is continuously updated during the run.
3. In each call, a line is appended for each layer (i.e. module) containing the named buffers.
"""
with ExitStack() as stack:
files = {}
writers = {}
for buf_name in buffer_names:
fname = self.get_fname(buf_name)
new = not os.path.isfile(fname)
files[buf_name] = stack.enter_context(open(fname, 'a'))
writer = csv.writer(files[buf_name])
if new:
writer.writerow(
['Layer', 'Epoch', 'Step', 'Total', 'Values'])
writers[buf_name] = writer
for n, m in model.named_modules():
for buffer_name in buffer_names:
try:
p = getattr(m, buffer_name)
except AttributeError:
continue
writer = writers[buffer_name]
if isinstance(p, (list, torch.nn.ParameterList)):
values = []
for v in p:
values += v.view(-1).tolist()
else:
values = p.view(-1).tolist()
writer.writerow([
utils.normalize_module_name(n) + '.' +
buffer_name, epoch, completed,
int(total)
] + values)
def make_fc(model, fc_module, g, name, seq_id, layer_id):
fc = SimpleNamespace()
fc.type = "Linear"
fc.name = name
fc.id = layer_id
fc.t = seq_id
# Use the SummaryGraph to obtain some other details of the models
fc_op = g.find_op(normalize_module_name(name))
assert fc_op is not None
fc.weights_vol = fc_op['attrs']['weights_vol']
fc.macs = fc_op['attrs']['MACs']
fc.n_ofm = fc_op['attrs']['n_ofm']
fc.n_ifm = fc_op['attrs']['n_ifm']
fc_pname = name + ".weight"
fc_p = utils.utils.model_find_param(model, fc_pname)
fc.ofm_h = g.param_shape(fc_op['outputs'][0])[0]
fc.ofm_w = g.param_shape(fc_op['outputs'][0])[1]
fc.ifm_h = g.param_shape(fc_op['inputs'][0])[0]
fc.ifm_w = g.param_shape(fc_op['inputs'][0])[1]
return fc
def make_conv(model, conv_module, g, name, seq_id, layer_id):
conv = SimpleNamespace()
conv.type = "Conv2D"
conv.name = name
conv.id = layer_id
conv.t = seq_id
conv.k = conv_module.kernel_size[0]
conv.stride = conv_module.stride
# Use the SummaryGraph to obtain some other details of the models
conv_op = g.find_op(normalize_module_name(name))
assert conv_op is not None
conv.weights_vol = conv_op['attrs']['weights_vol']
conv.macs = conv_op['attrs']['MACs']
conv.n_ofm = conv_op['attrs']['n_ofm']
conv.n_ifm = conv_op['attrs']['n_ifm']
conv_pname = name + ".weight"
conv_p = utils.utils.model_find_param(model, conv_pname)
conv.ofm_h = g.param_shape(conv_op['outputs'][0])[2]
conv.ofm_w = g.param_shape(conv_op['outputs'][0])[3]
conv.ifm_h = g.param_shape(conv_op['inputs'][0])[2]
conv.ifm_w = g.param_shape(conv_op['inputs'][0])[3]
return conv
def log_model_buffers(self, model, buffer_names, tag_prefix, epoch,
completed, total, freq):
"""Logs values of model buffers.
Notes:
1. Each buffer provided in 'buffer_names' is displayed in a separate table.
2. Within each table, each value is displayed in a separate column.
"""
datas = {name: [] for name in buffer_names}
maxlens = {name: 0 for name in buffer_names}
for n, m in model.named_modules():
for buffer_name in buffer_names:
try:
p = getattr(m, buffer_name)
except AttributeError:
continue
data = datas[buffer_name]
values = p if isinstance(
p,
(list, torch.nn.ParameterList)) else p.view(-1).tolist()
data.append([
utils.normalize_module_name(n) + '.' + buffer_name, *values
])
maxlens[buffer_name] = max(maxlens[buffer_name], len(values))
def find_op(self, lost_op_name):
return self.ops.get(utils.normalize_module_name(lost_op_name), None)