ede593804169a727e65b6723ec8a075afdf4c8c2,src/sdk/pynni/tests/test_compressor.py,CompressorTestCase,test_torch_l1filter_pruner,#CompressorTestCase#,166

Before Change

        model = TorchModel()
        config_list = [{"sparsity": 0.2, "op_types": ["Conv2d"], "op_names": ["conv1"]},
                       {"sparsity": 0.6, "op_types": ["Conv2d"], "op_names": ["conv2"]}]
        pruner = torch_compressor.L1FilterPruner(model, config_list)

        model.conv1.weight.data = torch.tensor(w).float()
        model.conv2.weight.data = torch.tensor(w).float()
        layer1 = torch_compressor.compressor.LayerInfo("conv1", model.conv1)
        mask1 = pruner.calc_mask(layer1, config_list[0], if_calculated=torch.tensor(0))
        layer2 = torch_compressor.compressor.LayerInfo("conv2", model.conv2)
        mask2 = pruner.calc_mask(layer2, config_list[1], if_calculated=torch.tensor(0))
        assert all(torch.sum(mask1["weight"], (1, 2, 3)).numpy() == np.array([0., 27., 27., 27., 27.]))
        assert all(torch.sum(mask2["weight"], (1, 2, 3)).numpy() == np.array([0., 0., 0., 27., 27.]))

    def test_torch_slim_pruner(self):

After Change

        w = np.array([np.zeros((3, 3, 3)), np.ones((3, 3, 3)), np.ones((3, 3, 3)) * 2,
                      np.ones((3, 3, 3)) * 3, np.ones((3, 3, 3)) * 4])
        model = TorchModel()
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
        config_list = [{"sparsity": 0.2, "op_types": ["Conv2d"], "op_names": ["conv1"]},
                       {"sparsity": 0.6, "op_types": ["Conv2d"], "op_names": ["conv2"]}]
        pruner = torch_compressor.L1FilterPruner(model, config_list, optimizer)

        model.conv1.module.weight.data = torch.tensor(w).float()
        model.conv2.module.weight.data = torch.tensor(w).float()
        mask1 = pruner.calc_mask(model.conv1)
        mask2 = pruner.calc_mask(model.conv2)
        assert all(torch.sum(mask1["weight_mask"], (1, 2, 3)).numpy() == np.array([0., 27., 27., 27., 27.]))
        assert all(torch.sum(mask2["weight_mask"], (1, 2, 3)).numpy() == np.array([0., 0., 0., 27., 27.]))

    def test_torch_slim_pruner(self):

Instances