Pruning the Convnets

evcu · February 12, 2018, 9:23pm

I would suggest using a wrapper module like the one here:

evcu/pytorchpruner/blob/76b3fb7ded9af7c6a84d4792efdaed235cb0dd06/pytorchpruner/modules.py#L4


import torch
from torch.nn.modules import Module


class MaskedModule(Module):
r"""Implements masked module for prunning etc...
it creates a mask for each layer and holds those masks inside a dictionary belongs to
this module
"""
DEFAULT_MASKED_MODULES = (torch.nn.Conv2d,torch.nn.Linear)
def __init__(self, module, masked_modules=DEFAULT_MASKED_MODULES):
    super(MaskedModule, self).__init__()
    self.module = module
    self._mask_dict = {} # nn.Module->torch.ByteTensor
    self._inp_dict = {} # nn.Module->torch.Tensor

The library itself is an ongoing project of mine about pruning and calculating saliencies.