def test_magnitude_pruning():
# Create a 4-D tensor of 1s
a = torch.ones(3, 64, 32, 32)
# Change one element
a[1, 4, 17, 31] = 0.2
# Create a masks dictionary and populate it with one ParameterMasker
zeros_mask_dict = {}
masker = distiller.ParameterMasker('a')
zeros_mask_dict['a'] = masker
# Try to use a MagnitudeParameterPruner with defining a default threshold
with pytest.raises(AssertionError):
pruner = distiller.pruning.MagnitudeParameterPruner("test", None)
# Now define the default threshold
thresholds = {"*": 0.4}
pruner = distiller.pruning.MagnitudeParameterPruner("test", thresholds)
assert distiller.sparsity(a) == 0
# Create a mask for parameter 'a'
pruner.set_param_mask(a, 'a', zeros_mask_dict, None)
assert common.almost_equal(distiller.sparsity(zeros_mask_dict['a'].mask), 1/distiller.volume(a))
# Let's now use the masker to prune a parameter
masker = zeros_mask_dict['a']
masker.apply_mask(a)
assert common.almost_equal(distiller.sparsity(a), 1/distiller.volume(a))
# We can use the masker on other tensors, if we want (and if they have the correct shape).
# Remember that the mask was created already, so we're not thresholding - we are pruning
b = torch.ones(3, 64, 32, 32)
b[:] = 0.3
masker.apply_mask(b)
assert common.almost_equal(distiller.sparsity(b), 1/distiller.volume(a))
def create_model_masks(model):
# Create the masks
zeros_mask_dict = {}
for name, param in model.named_parameters():
masker = distiller.ParameterMasker(name)
zeros_mask_dict[name] = masker
return zeros_mask_dict
def setup_test(arch, dataset, parallel):
model = create_model(False, dataset, arch, parallel=parallel)
assert model is not None
# Create the masks
zeros_mask_dict = {}
for name, param in model.named_parameters():
masker = distiller.ParameterMasker(name)
zeros_mask_dict[name] = masker
return model, zeros_mask_dict
