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!