Change shape of convolutional layers

1

Hey guys,

I’ve been trying out pytorch for a while and have a somewhat contrived used case: I am trying to change the shape of the weight tensor inside Conv2d layers in order to do some filter pruning on a pre-trained model:

I wrote some code to change the shape of all the conv layers in the model.

for layer in model.modules():
        if isinstance(layer, nn.Conv2d):
            print("Pruning Conv-{}".format(index))
            filters = layer.weight.data.numpy()
            new_filters = prune_layer(filters)  # reshape the tensor

            # Update layer with new tensor
            layer.weight.data = torch.from_numpy(new_filters).float()
            layer.out_channels = new_filters.shape[0]

The problem I have is that since the out_channels of the ith layer change, I also need to adjust the in_channels of the
ith+1 layer and from modules() I can’t guarantee that the order of the nn.Conv2d modules is the same as the one defined in the forward function, so I might end up changing the channels for the wrong layers.

Is there a better approach to this? Am I doing it completely wrong?

Thanks for the help! :slight_smile:

2

If you don’t have skip-layer connections, the order should usually be the same. Just note that you need to prune weight and bias in Conv2d and BatchNorm2d (if you use them).

If you have skip-layer connections, you can build a dependency graph about who passes output to who. You can then use the dependency graph to trace which convs take input from which convs. One example that builds such a graph can be found here (look up pytorch_visualize.py). This example performs graph visualization.

A much simpler way, however, is to “hijack” the forward function of Conv2d, BatchNorm2d, activations functions and so on, like so:

old_forwards = {nn.Conv2d: nn.Conv2d.forward}

def hijack_forward(self, input):
    if hasattr(input, 'drop_mask'):
        drop_input_channels(self, var.drop_mask)  # modify weights and biases
    drop_mask = drop_output_channels(self)        # modify weights and biases
    output = old_forwards[self.__class__](self, input)
    output.drop_mask = drop_mask
    return output

nn.Conv2d.forward = hijack_forward
# perform one forward pass
# reset hijacked forward functions
1 Like
3

Thanks for your answer!

The hijacking approach looks nice but I still need to keep track of dependencies between layers right?
In order to update the in_channels of some of them, I mean.

4

var.drop_mask is computed from a previous convolution layer. It tells you which input channels should be dropped.