test_pruner_torch.py 8.57 KB
Newer Older
1
2
3
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

4
import random
5
6
import unittest

7
import numpy
8
9
10
import torch
import torch.nn.functional as F

J-shang's avatar
J-shang committed
11
import nni
J-shang's avatar
J-shang committed
12
from nni.compression.pytorch.pruning import (
13
14
15
16
17
18
19
20
    LevelPruner,
    L1NormPruner,
    L2NormPruner,
    SlimPruner,
    FPGMPruner,
    ActivationAPoZRankPruner,
    ActivationMeanRankPruner,
    TaylorFOWeightPruner,
21
22
    ADMMPruner,
    MovementPruner
23
)
J-shang's avatar
J-shang committed
24
from nni.algorithms.compression.v2.pytorch.utils import compute_sparsity_mask2compact
25
26
27
28
29


class TorchModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
30
31
32
33
34
        self.conv1 = torch.nn.Conv2d(1, 10, 5, 1)
        self.bn1 = torch.nn.BatchNorm2d(10)
        self.conv2 = torch.nn.Conv2d(10, 20, 5, 1)
        self.bn2 = torch.nn.BatchNorm2d(20)
        self.fc1 = torch.nn.Linear(4 * 4 * 20, 100)
35
36
37
38
39
40
41
        self.fc2 = torch.nn.Linear(100, 10)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.max_pool2d(x, 2, 2)
42
        x = x.view(-1, 4 * 4 * 20)
43
44
45
46
47
48
49
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)


def trainer(model, optimizer, criterion):
    model.train()
50
51
52
53
54
55
56
57
    for _ in range(10):
        input = torch.rand(10, 1, 28, 28)
        label = torch.Tensor(list(range(10))).type(torch.LongTensor)
        optimizer.zero_grad()
        output = model(input)
        loss = criterion(output, label)
        loss.backward()
        optimizer.step()
58
59
60


def get_optimizer(model):
J-shang's avatar
J-shang committed
61
    return nni.trace(torch.optim.SGD)(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
62
63
64
65
66
67
68
69
70
71
72
73
74


criterion = torch.nn.CrossEntropyLoss()


class PrunerTestCase(unittest.TestCase):
    def test_level_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
        pruner = LevelPruner(model=model, config_list=config_list)
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
75
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
76

77
78
79
80
81
82
83
84
85
86
    def test_level_pruner_bank(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.7}]
        pruner = LevelPruner(model=model, config_list=config_list, mode='balance', balance_gran=[5])
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
        # round down cause to lower sparsity
        assert sparsity_list[0]['total_sparsity'] == 0.6

87
88
89
90
91
92
93
94
    def test_l1_norm_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
        pruner = L1NormPruner(model=model, config_list=config_list, mode='dependency_aware',
                              dummy_input=torch.rand(10, 1, 28, 28))
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
95
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
96
97
98
99
100
101
102
103
104

    def test_l2_norm_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
        pruner = L2NormPruner(model=model, config_list=config_list, mode='dependency_aware',
                              dummy_input=torch.rand(10, 1, 28, 28))
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
105
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
106
107
108
109
110
111
112
113
114

    def test_fpgm_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
        pruner = FPGMPruner(model=model, config_list=config_list, mode='dependency_aware',
                            dummy_input=torch.rand(10, 1, 28, 28))
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
115
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
116
117
118
119

    def test_slim_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['BatchNorm2d'], 'total_sparsity': 0.8}]
120
        pruner = SlimPruner(model=model, config_list=config_list, trainer=trainer, traced_optimizer=get_optimizer(model),
121
122
123
124
                            criterion=criterion, training_epochs=1, scale=0.001, mode='global')
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
125
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
126
127
128
129
130

    def test_activation_mean_rank_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
        pruner = ActivationMeanRankPruner(model=model, config_list=config_list, trainer=trainer,
131
                                          traced_optimizer=get_optimizer(model), criterion=criterion, training_batches=5,
132
133
134
135
136
                                          activation='relu', mode='dependency_aware',
                                          dummy_input=torch.rand(10, 1, 28, 28))
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
137
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
138
139
140
141
142

    def test_taylor_fo_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
        pruner = TaylorFOWeightPruner(model=model, config_list=config_list, trainer=trainer,
143
                                      traced_optimizer=get_optimizer(model), criterion=criterion, training_batches=5,
144
145
146
147
                                      mode='dependency_aware', dummy_input=torch.rand(10, 1, 28, 28))
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
148
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82
149
150
151
152

    def test_admm_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8, 'rho': 1e-3}]
153
        pruner = ADMMPruner(model=model, config_list=config_list, trainer=trainer, traced_optimizer=get_optimizer(model),
154
155
156
157
                            criterion=criterion, iterations=2, training_epochs=1)
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
158
159
160
161
162
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82

    def test_movement_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
163
        pruner = MovementPruner(model=model, config_list=config_list, trainer=trainer, traced_optimizer=get_optimizer(model),
164
165
166
167
168
169
                                criterion=criterion, training_epochs=5, warm_up_step=0, cool_down_beginning_step=4)
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82

170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
class FixSeedPrunerTestCase(unittest.TestCase):
    def test_activation_apoz_rank_pruner(self):
        model = TorchModel()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
        pruner = ActivationAPoZRankPruner(model=model, config_list=config_list, trainer=trainer,
                                            traced_optimizer=get_optimizer(model), criterion=criterion, training_batches=5,
                                            activation='relu', mode='dependency_aware',
                                            dummy_input=torch.rand(10, 1, 28, 28))
        pruned_model, masks = pruner.compress()
        pruner._unwrap_model()
        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
        assert 0.78 < sparsity_list[0]['total_sparsity'] < 0.82

    def setUp(self) -> None:
        # fix seed in order to solve the random failure of ut
        random.seed(1024)
        numpy.random.seed(1024)
        torch.manual_seed(1024)

    def tearDown(self) -> None:
        # reset seed
        import time
        now = int(time.time() * 100)
        random.seed(now)
        seed = random.randint(0, 2 ** 32 - 1)
        random.seed(seed)
        numpy.random.seed(seed)
        torch.manual_seed(seed)
198
199
200

if __name__ == '__main__':
    unittest.main()