microsoft / nni

An open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning.
https://nni.readthedocs.io
MIT License
13.99k stars 1.81k forks source link

Does nni contain some new pruning base methods? #4567

Open typhoonlee opened 2 years ago

typhoonlee commented 2 years ago

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:

How to reproduce it?:

J-shang commented 2 years ago

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. https://github.com/microsoft/nni/blob/632df3ea2f99f3c8e4d2a16fab6ebe4303609da1/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py#L505

you need to write a new collector to collect layers' output, https://github.com/microsoft/nni/blob/632df3ea2f99f3c8e4d2a16fab6ebe4303609da1/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py#L590

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. https://github.com/microsoft/nni/blob/632df3ea2f99f3c8e4d2a16fab6ebe4303609da1/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py#L26

You can see how pruner.compress() works in: https://github.com/microsoft/nni/blob/632df3ea2f99f3c8e4d2a16fab6ebe4303609da1/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py#L106

J-shang commented 2 years ago

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

typhoonlee commented 2 years ago

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:blush:, thank you very much for your patience!

typhoonlee commented 2 years ago

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: 0222_error1 0222_error2 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 commented 2 years ago

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 commented 2 years ago

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.