def test_get_weight_parameter(self):
""" Check whether parameters can be get from Module. """
self.assertIsNotNone(
model_utils.get_weight_parameter(MaskConv2d(32, 32, 3)))
weight_groups = model_utils.get_weight_parameter(
GroupConv2d(32, 64, 3, groups=2))
self.assertIsNotNone(weight_groups)
self.assertIsInstance(weight_groups, torch.Tensor)
self.assertEqual(weight_groups.shape[0], 64)
self.assertEqual(weight_groups.shape[1], 16)
def find_groupable_modules(self, model, G=1, MCPG=None, g_cfg=None):
""" Find modules that can be turned into GroupConv2d. """
mods = []
for name, mod in model.named_modules():
if isinstance(mod, MaskConv2d):
# we assume that if the number of groups is not divisible
# it is not grouped
# TODO: support hybrid group size
weight = model_utils.get_weight_parameter(mod)
F, C = weight.shape[:2]
if not GroupConv2d.groupable(
C, F, groups=G, max_channels_per_group=MCPG):
continue
# get the desired group number
if g_cfg and name in g_cfg:
G_ = g_cfg[name]["G"]
else:
G_ = GroupConv2d.get_num_groups(
C, F, groups=G, max_channels_per_group=MCPG)
# we also need to check whether GSP conditions are met.
if not model_utils.is_gsp_satisfied(mod, G_):
continue
mods.append((name, mod))
return mods
def get_num_groups(name, mod):
G_ = G # choose G in the beginning
W = model_utils.get_weight_parameter(mod)
F, C = W.shape[:2]
# how to override G_
if g_cfg is not None:
if name in g_cfg:
G_ = g_cfg[name]["G"]
# do some verification
assert F == g_cfg[name]["F"] and C == g_cfg[name]["C"]
else:
G_ = 1 # HACK - we don't want to have G=0 in further processing
elif MCPG > 0:
if GroupConv2d.groupable(C, F, max_channels_per_group=MCPG):
G_ = GroupConv2d.get_num_groups(C,
F,
max_channels_per_group=MCPG)
else:
logging.warn(
"Module {} is not groupable under MCPG={}, set its G to 1".
format(name, MCPG))
G_ = 1
return G_
def prune_module(self,
name,
mod,
*args,
g_cfg=None,
fake_mask=False,
**kwargs):
""" Prune a single module.
We expect that after pruning, the mask in mod can be
updated to a pruned result.
Args:
name(str): name of the module
mod(MaskConv2d): the module to be pruned.
"""
# TODO maybe not pass by kwargs
G = self.args.num_groups
if g_cfg is not None and name in g_cfg:
G = g_cfg[name]["G"]
# do some verification
W = model_utils.get_weight_parameter(mod)
F, C = W.shape[:2]
assert F == g_cfg[name]["F"] and C == g_cfg[name]["C"]
if fake_mask and isinstance(mod, MaskConv2d):
mod.fake_mask = True
prune_utils.prune_module(mod, G=G, MCPG=self.args.mcpg, **kwargs)
def get_num_groups(name, mod):
G_ = G # choose G in the beginning
W = model_utils.get_weight_parameter(mod)
F, C = W.shape[:2]
if not data_parallel:
name = 'module.' + name
# how to override G_
if g_cfg is not None:
if name in g_cfg:
G_ = g_cfg[name]['G']
# do some verification
if G_ != 1: # HACK
assert F == g_cfg[name][
'F'], 'F={} does not match cfg={}'.format(
F, g_cfg[name]['F'])
assert C == g_cfg[name][
'C'], 'C={} does not match cfg={}'.format(
C, g_cfg[name]['C'])
else:
G_ = 1 # HACK - we don't want to have G=0 in further processing
elif MCPG > 0:
if GroupConv2d.groupable(C, F, max_channels_per_group=MCPG):
G_ = GroupConv2d.get_num_groups(C,
F,
max_channels_per_group=MCPG)
else:
logging.warn(
'Module {} is not groupable under MCPG={}, set its G to 1'.
format(name, MCPG))
G_ = 1
return G_
def export(self):
""" Export function """
# load the original model
model = self.load_model(self.args)
self.evaluate(model)
# all these modules can be further converted to GroupConv2d
G = self.args.num_groups
MCPG = self.args.mcpg
g_cfg = self.load_group_cfg(self.args)
logging.info("Finding all groupable modules ...")
mods = self.find_groupable_modules(model, G=G, MCPG=MCPG, g_cfg=g_cfg)
logging.info("Found {} modules".format(len(mods)))
# for each module, create corresponding group convolution
# parameters, and use them to update the state_dict
logging.info("Generating GroupConv2d parameters ...")
state_dict = torch.load(self.args.resume)["state_dict"]
for name, mod in mods:
weight = model_utils.get_weight_parameter(mod)
F, C = weight.shape[:2]
# will update G correspondingly
# TODO put this logic somewhere else, frequently reused
if g_cfg and name in g_cfg:
G_ = g_cfg[name]["G"]
else:
G_ = GroupConv2d.get_num_groups(C, F, MCPG, groups=G)
wg, ind_in, ind_out = model_utils.get_group_parameters(mod, G_)
# update the state_dict
del state_dict[name + ".mask"] # delete mask key
# purge weight in MaskConv2d
if (name + ".weight") in state_dict:
del state_dict[name + ".weight"]
state_dict[name + ".conv2d.weight"] = torch.from_numpy(wg)
state_dict[name + ".ind_in"] = torch.from_numpy(ind_in).long()
state_dict[name + ".ind_out"] = torch.from_numpy(ind_out).long()
# create model
# leave indices to None, will set up when loading state_dict
model = self.create_model(self.args,
groups=G,
max_channels_per_group=MCPG,
mask=False)
if g_cfg: # should post update
self.update_model_by_group_cfg(model, g_cfg)
# insert weights
model.load_state_dict(state_dict)
self.evaluate(model)
def find_groupable_modules(self, model):
""" Find modules that can be grouped. """
mods = [] # return list
for name, mod in model.named_modules():
if isinstance(mod, MaskConv2d):
W = model_utils.get_weight_parameter(mod)
F, C = W.shape[:2]
ff, fc = factors(F), factors(C)
if len(ff.intersection(fc)) <= 1:
continue
mods.append((name, mod))
return mods
def prune_module(self, name, mod, G, **kwargs):
""" Prune a specific module.
NOTE: G is known at this moment.
"""
assert isinstance(mod, MaskConv2d)
assert G >= 1, "{} has G={} smaller than 1".format(name, G)
W = model_utils.get_weight_parameter(mod)
C_out, C_in = W.shape[:2]
if G == 1 or (C_out % G != 0 or C_in % G != 0):
# NOTE: we return if this module cannot be pruned
return
prune_utils.prune_module(mod, G=G, **kwargs)
def find_group_candidates(self, mod, **kwargs):
""" Find group number candidates in module.
Note: use kwargs to pass additional requirements.
"""
W = model_utils.get_weight_parameter(mod)
F, C = W.shape[:2]
# common divisors
Gs = list(sorted(factors(F).intersection(factors(C))))
del Gs[0] # should be 1
costs = []
for G in Gs:
_, _, cost = mask_utils.run_mbm(W, G)
costs.append(cost)
return Gs, costs
def run_opt(self):
""" Run the actual optimization """
logging.info('Finding the optimal group configuration ...')
model = self.load_model(self.args)
g_conf = OrderedDict()
for idx, (name, mod) in enumerate(self.find_groupable_modules(model)):
# the sequence here is important
Gs, costs = list(self.find_group_candidates(mod))
W = model_utils.get_weight_parameter(mod)
F, C = W.shape[:2]
# TODO: simply select the one with the most cost
G = Gs[np.argmax(costs)]
# print(np.max(costs) / (F * C))
# update the dictionary
g_conf[name] = {'id': idx, 'F': F, 'C': C, 'G': G}
return g_conf
