In a recent paper https://arxiv.org/abs/2003.11883 3 the authors are talking about forwarding only 25% of random sampled channels of a featuremap through a conv-layer while “bypassing” the others to reduce the memory footprint.
Here is the corresponding section in that paper:

image755×430 76.5 KB

How would you achieve something like this in PyTorch?
I figured you could achieve this with something like in the following example, but maybe there is another simpler and obvious way that I am missing?
class SampleConv(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
        super().__init__(
            in_channels, out_channels,
            kernel_size, stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias)

        self.channel_sample_ratio = 0.25

    def forward(self, input):
        sampled_channels = torch.rand(x.size(1), device=x.device, requires_grad=False)

        weight = self.weight[:, sampled_channels < self.channel_sample_ratio]
        bias = self.bias

        y = F.conv2d(input[:, sampled_channels < self.channel_sample_ratio], weight, bias, self.stride, self.padding, self.dilation, self.groups)
        return y

Also I measured the memory consumption of above SampleConv and its regular nn.Conv2d-equivalent and it doen’t seem like that it requires 75% less memory.