Does nni contain some new pruning base methods？

[typhoonlee]

Describe the issue:
Sorry to bother you, in nni, the current basic pruning strategies are Norm, FPGM, ActivationPruner, TaylorFOWeightPruner, etc. Are there other basic strategies, such as HRank, if I want to add it myself, where do I need to rewrite it?

Environment:

• NNI version: V2.6
• Training service (local|remote|pai|aml|etc): local
• Client OS: linux
• Server OS (for remote mode only):
• Python version: 3.8
• PyTorch/TensorFlow version: pytorch1.8
• Is conda/virtualenv/venv used?: conda
• Is running in Docker?: no

How to reproduce it?:

[J-shang]

Wow, it is a good question, let's use ActivationPruner as an example, I didn't read HRank paper carefully, but if I understand correctly, HRank and APoZ have a similar pruning process.

nni/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py

Line 505 in 632df3e

class ActivationPruner(BasicPruner):

you need to write a new collector to collect layers' output,

nni/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py

Line 590 in 632df3e

def _collector(self, buffer: List) -> Callable[[Module, Tensor, Tensor], None]:

SingleHookTrainerBasedDataCollector.collect() will return the dict {op_names: buffer}

Then you need to write a hrank metric calculator to use this returned buffer dict to calculate the HRank. It is similar to NormMetricsCalculator, NormMetricsCalculator calculate norm on keeped_dim, HRankMetricsCalculator calculate HRank on keeped_dim.

nni/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py

Line 26 in 632df3e

class NormMetricsCalculator(MetricsCalculator):

You can see how pruner.compress() works in:

nni/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py

Line 106 in 632df3e

def compress(self) -> Tuple[Module, Dict]:

[J-shang]

Welcome to contribute to NNI if you write this new Pruner 😉 and feel free to contact us if you meet some issues.

[typhoonlee]

Welcome to contribute to NNI if you write this new Pruner 😉 and feel free to contact us if you meet some issues.

I will try to rewrite it first😊, thank you very much for your patience！

[typhoonlee]

I modified nni/nni/algorithms/compression/v2/pytorch/pruning/tools/metrix_caculator.py as this:

class MeanRankMetricsCalculator(MetricsCalculator):
    """
    The data value format is a two-element list [batch_number, batch_wise_activation_sum].
    This metric simply calculate the average on `self.dim`, then divide by the batch_number.
    MeanRank pruner uses this to calculate metric.
    """
    def calculate_metrics(self, data: Dict[str, List[Tensor]]) -> Dict[str, Tensor]:
        metrics = {}
        for name, (num, activation_sum) in data.items():
            keeped_dim = list(range(len(activation_sum.size()))) if self.dim is None else self.dim
            across_dim = list(range(len(activation_sum.size())))
            [across_dim.pop(i) for i in reversed(keeped_dim)]
            # metrics[name] = torch.mean(activation_sum, across_dim) / num

            # modified on 220222: Implementing the HRank Strategy
            activation_sum_rank = torch.matrix_rank(activation_sum)
            activation_sum_rank1 = activation_sum_rank.float()
            metrics[name] = torch.mean(activation_sum_rank1, axis=0) / num
        return metrics

The rest is consistent with ActivationMeanRankPruner. But got this error：


But I obviously only pruned the filter dimension, and the pruning rate was 0.125：config_list = [{'op_types': ['Conv2d'], 'sparsity_per_layer': 0.125}, {'exclude': True, 'op_names': ['attconv']}]

[J-shang]

you could check the values in metrics, I think there may have some problems.

I don't know your whole code, but if the data you collected is [batch_num, output_dim, feature_map_dims...]

metrics[name] = torch.mean(activation_sum_rank1, axis=across_dim) / num

When you initialize this MeanRankMetricsCalculator, set dim=1:

MeanRankMetricsCalculator(dim=1)

[typhoonlee]

you could check the values in metrics, I think there may have some problems.

I don't know your whole code, but if the data you collected is [batch_num, output_dim, feature_map_dims...]

metrics[name] = torch.mean(activation_sum_rank1, axis=across_dim) / num

When you initialize this MeanRankMetricsCalculator, set dim=1:

MeanRankMetricsCalculator(dim=1)

Although it is not clear why，I solved this problem by setting num to 1.