Recently I’m doing some research on network pruning and built a custom layer for my vgg network.
My custom layer
Mask contains a set of parameters that have the same shape as the input feature map. However, after I implemented the module together with my other conv layers using
nn.Sequential, the parameters
self.mask does not get updated.
I think perhaps my custom module broke some rules in Pytorch, can you please give me some advice on getting it work?
My custom layer is defined as following:
class Mask(nn.Module): def __init__(self,bits=1): super(Mask, self).__init__() self.mask = None self.bits = bits self.weight_shape = None self.channel = None self.height = None self.width = None def forward(self,input): if type(self.mask) == type(None): self.mask = Parameter(torch.Tensor(input.shape[1:]).cuda()) nn.init.constant_(self.mask,val=1.0) self.weight_shape = input.shape[1:] self.channel = input.shape self.height = input.shape self.width = input.shape result = input*self.mask return result
Mask layers are built using:
def make_layers(cfg, batch_norm=False): layers =  in_channels = 3 for v in cfg: if v == 'M': layers += [nn.MaxPool2d(kernel_size=2, stride=2)] else: mask = utils.Mask(bits=1).cuda() conv2d = utils.Conv2d(in_channels, v, kernel_size=3, padding=1) if batch_norm: layers += [mask,conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] else: layers += [mask,conv2d, nn.ReLU(inplace=True)] in_channels = v return nn.Sequential(*layers)
I initially defined the
None and initialized it in the
forward() function because my intention was to mask the input of the layer, so I cannot fetch the size of input during initializing the module. But I suspect there might be something wrong here. If so, is there any other way I can do it correctly? Thanks in advance!