def test_mask_squash(self): model = Model() sparsifier = WeightNormSparsifier() sparsifier.prepare(model, config=None) sparsifier.squash_mask() for g in sparsifier.module_groups: module = g['module'] assert not is_parametrized(module, 'weight') assert not hasattr(module, 'mask')
def test_sparsity_levels(self): sparsity_levels = [-1.0, 0.0, 0.5, 1.0, 2.0] sparse_block_shapes = [(1, 1), (1, 4), (2, 2), (4, 1)] zeros_per_blocks = [0, 1, 2, 3, 4] testcases = itertools.tee(itertools.product(sparsity_levels, sparse_block_shapes, zeros_per_blocks)) # Create a config and model with all the testcases model = nn.Sequential() sparsifier = WeightNormSparsifier() sparsity_per_layer_config = [] p = re.compile(r'[-\.\s]') for sl, sbs, zpb in testcases[0]: # Make sure the number of zeros is not > values in a block if zpb > sbs[0] * sbs[1]: continue layer_name = f'{sl}_{sbs}_{zpb}' layer_name = p.sub('_', layer_name) layer = nn.Linear(12, 12, bias=False) layer.weight = nn.Parameter(torch.ones(12, 12)) model.add_module(layer_name, layer) config = { 'fqn': layer_name, 'sparsity_level': sl, 'sparse_block_shape': sbs, 'zeros_per_block': zpb } sparsity_per_layer_config.append(config) sparsifier.prepare(model, sparsity_per_layer_config) sparsifier.step() sparsifier.squash_mask() model.eval() for sl, sbs, zpb in testcases[1]: if zpb > sbs[0] * sbs[1]: continue layer_name = f'{sl}_{sbs}_{zpb}' layer_name = p.sub('_', layer_name) layer = getattr(model, layer_name) # Level of sparsity is achieved sparse_mask = (layer.weight == 0).float() if zpb == 0: assert sparse_mask.mean() == 0 else: # Ratio of individual zeros in the tensor true_sl = min(max(sl, 0.0), 1.0) true_sl = true_sl * zpb / sbs[0] / sbs[1] assert sparse_mask.mean() == true_sl