Pruning: How to preserve the number of output channels of a particular layer?

[saravanabalagi]

I have a simple conv model and am able to obtain pruning masks and use ModelSpeedup to alter the model successfully.
However, I am unable to preserve the output channels of the last layer.

Minimal Code

# %%
import torch
import torch.nn as nn

from nni.compression.pruning import L1NormPruner
from nni.compression.utils import auto_set_denpendency_group_ids
from nni.compression.speedup import ModelSpeedup

# %%
class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 40, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(40)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(40, 80, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(80)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        return x
    
model = ConvNet()
num_params_unpruned = sum(p.numel() for p in model.parameters())
dummy_input = torch.randn(1, 3, 32, 32)
dummy_output = model(dummy_input)
print(dummy_output.shape)

# %%
sparsity_ratio = 0.5
config_list = [{
    'op_types': ['Conv2d'],
    'sparse_ratio': sparsity_ratio,
}]
config_list = auto_set_denpendency_group_ids(model, config_list, [dummy_input])
pruner = L1NormPruner(model, config_list)
_, masks = pruner.compress()
pruner.unwrap_model()
model = ModelSpeedup(model, [dummy_input], masks, batch_size=1, garbage_collect_values=False).speedup_model()

# %%
num_params_pruned = sum(p.numel() for p in model.parameters())
print(f'Number of parameters before pruning: {num_params_unpruned}')
print(f'Number of parameters after pruning: {num_params_pruned}')

num_params_diff = num_params_unpruned - num_params_pruned
prune_ratio = num_params_diff / num_params_unpruned
print(f'Number of parameters pruned: {num_params_diff}')
print(f'Parameter ratio: {(1-prune_ratio)*100:.2f}%')

I tried adding the layer names in exclude_op_names in config, but it does not seem to respect it. I even tried adding the entire list and still it does as if this prop is not there.

sparsity_ratio = 0.5
config_list = [{
    'op_types': ['Conv2d'],
    'sparse_ratio': sparsity_ratio,
    'exclude_op_names': [
        'conv1.weight',
        'conv1.bias',
        'bn1.weight',
        'bn1.bias',
        'conv2.weight',
        'conv2.bias',
        'bn2.weight',
        'bn2.bias'
    ],
}]

Even if it works, I am no entirely if the pipeline will entirely skip the weights when mentioned. For example in the above model I would like to preserve the output channels of conv2.

I tried adding conv2* to exclude_op_names_re, then I get WARNING: no multi-dimension masks found. followed by a IndexError: list index out of range. It is understandable that it is trying to skip it, but I only want it to preserve the dim0 of the weights shape in conv2. I also tried adding bn2*, but it simply ignores it.

What is the correct way to preserve or freeze output channels of a particular layer in the model?

[saravanabalagi]

Found out that the op names do not and should not include weight or bias. So with that, using adding the last layer's name in the exclude_op_names just works:

config_list = [{
    'op_types': ['Conv2d'],
    'sparse_ratio': sparsity_ratio,
    'exclude_op_names': [
        'conv2',
    ]
}]

Log

Ouput shape: torch.Size([1, 80, 32, 32])
[2024-01-18 16:35:00] Start to speedup the model...
[2024-01-18 16:35:00] Resolve the mask conflict before mask propagate...
[2024-01-18 16:35:00] dim0 sparsity: 0.489796
[2024-01-18 16:35:00] dim1 sparsity: 0.000000
0 Filter
[2024-01-18 16:35:00] dim0 sparsity: 0.489796
[2024-01-18 16:35:00] dim1 sparsity: 0.000000
[2024-01-18 16:35:00] Infer module masks...
[2024-01-18 16:35:00] Propagate original variables
[2024-01-18 16:35:00] Propagate variables for placeholder: x, output mask:  0.0000 
[2024-01-18 16:35:00] Propagate variables for call_module: conv1, weight:  0.4898 bias:  0.4898 , output mask:  0.0000 
[2024-01-18 16:35:00] Propagate variables for call_module: bn1, , output mask:  0.0000 
[2024-01-18 16:35:00] Propagate variables for call_module: relu, , output mask:  0.0000 
[2024-01-18 16:35:01] Propagate variables for call_module: conv2, , output mask:  0.0000 
[2024-01-18 16:35:01] Propagate variables for call_module: bn2, , output mask:  0.0000 
[2024-01-18 16:35:01] Propagate variables for output: output, output mask:  0.0000 
[2024-01-18 16:35:01] Update direct sparsity...
[2024-01-18 16:35:01] Update direct mask for placeholder: x, output mask:  0.0000 
[2024-01-18 16:35:01] Update direct mask for call_module: conv1, weight:  0.4898 bias:  0.4898 , output mask:  0.4898 
[2024-01-18 16:35:02] Update direct mask for call_module: bn1, , output mask:  0.4898 
[2024-01-18 16:35:02] Update direct mask for call_module: relu, , output mask:  0.4898 
[2024-01-18 16:35:02] Update direct mask for call_module: conv2, , output mask:  0.0000 
[2024-01-18 16:35:02] Update direct mask for call_module: bn2, , output mask:  0.0000 
[2024-01-18 16:35:02] Update direct mask for output: output, output mask:  0.0000 
[2024-01-18 16:35:02] Update indirect sparsity...
[2024-01-18 16:35:02] Update indirect mask for output: output, output mask:  0.0000 
[2024-01-18 16:35:03] Update indirect mask for call_module: bn2, , output mask:  0.0000 
[2024-01-18 16:35:03] Update indirect mask for call_module: conv2, , output mask:  0.0000 
[2024-01-18 16:35:04] Update indirect mask for call_module: relu, , output mask:  0.4898 
[2024-01-18 16:35:04] Update indirect mask for call_module: bn1, , output mask:  0.4898 
[2024-01-18 16:35:04] Update indirect mask for call_module: conv1, weight:  0.4898 bias:  0.4898 , output mask:  0.4898 
[2024-01-18 16:35:04] Update indirect mask for placeholder: x, output mask:  0.0000 
[2024-01-18 16:35:04] Resolve the mask conflict after mask propagate...
[2024-01-18 16:35:04] dim0 sparsity: 0.489796
[2024-01-18 16:35:04] dim1 sparsity: 0.000000
0 Filter
[2024-01-18 16:35:04] dim0 sparsity: 0.489796
[2024-01-18 16:35:04] dim1 sparsity: 0.000000
[2024-01-18 16:35:04] Replace compressed modules...
[2024-01-18 16:35:04] replace module (name: conv1, op_type: Conv2d)
[2024-01-18 16:35:04] replace conv2d with in_channels: 3, out_channels: 25
[2024-01-18 16:35:04] replace module (name: bn1, op_type: BatchNorm2d)
[2024-01-18 16:35:04] replace batchnorm2d with num_features: 25
[2024-01-18 16:35:04] replace module (name: relu, op_type: ReLU)
[2024-01-18 16:35:04] replace module (name: conv2, op_type: Conv2d)
[2024-01-18 16:35:04] replace conv2d with in_channels: 25, out_channels: 80
[2024-01-18 16:35:04] replace module (name: bn2, op_type: BatchNorm2d)
[2024-01-18 16:35:04] replace batchnorm2d with num_features: 80
[2024-01-18 16:35:04] Speedup done.
Number of parameters before pruning: 36990
Number of parameters after pruning: 18990
Number of parameters pruned: 18000
Parameter ratio: 51.34%
ConvNet(
  (relu): ReLU(inplace=True)
  (conv1): Conv2d(3, 25, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (bn1): BatchNorm2d(25, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (conv2): Conv2d(25, 80, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (bn2): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)

However, using re with 'exclude_op_names_re': ['conv2*'] (or con*2 for example) does not work.

Perhaps for a more direct way to add in config to only do pruning on input channels, there is a granularity prop. It must be specified along with op_names and sparse_ratio in a separate config added to the config list.

config_list = [
    {
        'op_names': ['conv2'],
        'sparse_ratio': sparsity_ratio,
        'granularity': 'in_channel',
    },
    {
        'op_types': ['Conv2d'],
        'sparse_ratio': sparsity_ratio,
    }, 
]

There are a few warnings running it with this config

Log

Ouput shape: torch.Size([1, 40, 32, 32])
[2024-01-18 16:32:58] WARNING: bias have already configured, the new config will be ignored.
[2024-01-18 16:32:58] WARNING: weight have already configured, the new config will be ignored.
[2024-01-18 16:32:58] Start to speedup the model...
[2024-01-18 16:32:58] Resolve the mask conflict before mask propagate...
[2024-01-18 16:32:58] dim0 sparsity: 0.166667
[2024-01-18 16:32:58] dim1 sparsity: 0.434783
[2024-01-18 16:32:58] WARNING: both dim0 and dim1 masks found.
1 Filter
[2024-01-18 16:32:58] dim0 sparsity: 0.166667
[2024-01-18 16:32:58] dim1 sparsity: 0.434783
[2024-01-18 16:32:58] WARNING: both dim0 and dim1 masks found.
[2024-01-18 16:32:58] Infer module masks...
[2024-01-18 16:32:58] Propagate original variables
[2024-01-18 16:32:58] Propagate variables for placeholder: x, output mask:  0.0000 
[2024-01-18 16:32:59] Propagate variables for call_module: conv1, weight:  0.5000 bias:  0.5000 , output mask:  0.0000 
[2024-01-18 16:32:59] Propagate variables for call_module: bn1, , output mask:  0.0000 
[2024-01-18 16:32:59] Propagate variables for call_module: relu, , output mask:  0.0000 
[2024-01-18 16:32:59] Propagate variables for call_module: conv2, weight:  0.5000 bias:  0.0000 , output mask:  0.0000 
[2024-01-18 16:32:59] Propagate variables for call_module: bn2, , output mask:  0.0000 
[2024-01-18 16:32:59] Propagate variables for output: output, output mask:  0.0000 
[2024-01-18 16:32:59] Update direct sparsity...
[2024-01-18 16:33:00] Update direct mask for placeholder: x, output mask:  0.0000 
[2024-01-18 16:33:00] Update direct mask for call_module: conv1, weight:  0.5000 bias:  0.5000 , output mask:  0.5000 
[2024-01-18 16:33:00] Update direct mask for call_module: bn1, , output mask:  0.5000 
[2024-01-18 16:33:00] Update direct mask for call_module: relu, , output mask:  0.5000 
[2024-01-18 16:33:00] Update direct mask for call_module: conv2, weight:  0.5000 bias:  0.0000 , output mask:  0.0000 
[2024-01-18 16:33:00] Update direct mask for call_module: bn2, , output mask:  0.0000 
[2024-01-18 16:33:00] Update direct mask for output: output, output mask:  0.0000 
[2024-01-18 16:33:00] Update indirect sparsity...
[2024-01-18 16:33:01] Update indirect mask for output: output, output mask:  0.0000 
[2024-01-18 16:33:01] Update indirect mask for call_module: bn2, , output mask:  0.0000 
[2024-01-18 16:33:01] Update indirect mask for call_module: conv2, weight:  0.7000 bias:  0.0000 , output mask:  0.0000 
[2024-01-18 16:33:02] Update indirect mask for call_module: relu, , output mask:  0.7000 
[2024-01-18 16:33:02] Update indirect mask for call_module: bn1, , output mask:  0.7000 
[2024-01-18 16:33:02] Update indirect mask for call_module: conv1, weight:  0.5000 bias:  0.5000 , output mask:  0.7000 
[2024-01-18 16:33:02] Update indirect mask for placeholder: x, output mask:  0.0000 
[2024-01-18 16:33:02] Resolve the mask conflict after mask propagate...
[2024-01-18 16:33:02] dim0 sparsity: 0.166667
[2024-01-18 16:33:02] dim1 sparsity: 0.608696
[2024-01-18 16:33:02] WARNING: both dim0 and dim1 masks found.
1 Filter
[2024-01-18 16:33:02] dim0 sparsity: 0.166667
[2024-01-18 16:33:02] dim1 sparsity: 0.608696
[2024-01-18 16:33:02] WARNING: both dim0 and dim1 masks found.
[2024-01-18 16:33:02] Replace compressed modules...
[2024-01-18 16:33:02] replace module (name: conv1, op_type: Conv2d)
[2024-01-18 16:33:02] replace conv2d with in_channels: 3, out_channels: 6
[2024-01-18 16:33:02] replace module (name: bn1, op_type: BatchNorm2d)
[2024-01-18 16:33:03] replace batchnorm2d with num_features: 6
[2024-01-18 16:33:03] replace module (name: relu, op_type: ReLU)
[2024-01-18 16:33:03] replace module (name: conv2, op_type: Conv2d)
[2024-01-18 16:33:03] replace conv2d with in_channels: 6, out_channels: 40
[2024-01-18 16:33:03] replace module (name: bn2, op_type: BatchNorm2d)
[2024-01-18 16:33:03] replace batchnorm2d with num_features: 40
[2024-01-18 16:33:03] Speedup done.
Number of parameters before pruning: 7920
Number of parameters after pruning: 2460
Number of parameters pruned: 5460
Parameter ratio: 31.06%
ConvNet(
  (relu): ReLU(inplace=True)
  (conv1): Conv2d(3, 6, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (bn1): BatchNorm2d(6, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (conv2): Conv2d(6, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (bn2): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

The warnings mention that the weights and biases configured for more than once are ignored, the ones from the first config are kept. When swapping the order granularity config is simply ignored because of the same.

That said, it would be more intuitive to define the base case first and override it with a special case.