Batch Convolutional Layers - Similar to torch.bmm but for convolutional operators

[AntreasAntoniou]

🚀 Feature

Enable PyTorch to compute convolutional features given:

• A collection of batches of images of shape (b_i, b_j, c, h, w)
• A batch of kernel weights of shape (b_i, out_features, in_features, kernel_height, kernel_width)

Each batch in a collection is convolved only with its respective parameters. This can be currently achieved with the existing PyTorch operations but is rather slow compared to what a CUDA implementation could do.

Motivation

Meta-learning is slowly becoming a very popular framework for tackling problems. In almost every instance of meta-learning, there exists an inner loop optimization process, where a network, is updated, multiple times, given a task. There exists a batch of tasks, so currently, one has to process each task sequentially. This is necessary because the weights of a network change within the inner loop optimization process, thus doing this in parallel would require a convolution operator that can receive a batch of parameter kernels to be convolved with a batch of tasks.

Furthermore, genetic algorithms, and RL, often require similar operations, where each individual in a population changes it's inner state with each update, hence, requiring a similar mechanism.

Finally, in cases where Hyper Networks are used, if one chose to have a sample conditional weight matrix for a particular operation, then a Batch Convolutional operator would be necessary.

I think it's in PyTorch's best interests to have such a feature available, to get ahead of the curve.

Pitch

I have already implemented a variant of what I am proposing. When I am using a large collection size, then my method out performs sequential convolutions by a factor of 4x. However, when my collection size is close to 1, hence, effectively being a normal convolutional layer, my method is 20x slower than the default CUDA-based convolutional implementation. I need help by someone who can either write such CUDA-based implementations, or point me to the right direction so I can do it.

Find my existing implementation on this attached

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class BatchConvLayer(nn.Module):
    def __init__(self, input_shape, kernel_size, stride, padding, dilation, groups):
        super(BatchConvLayer, self).__init__()
        self.input_shape = input_shape
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        self.build_block()

    def image_to_patch_coordinates(self, input_image, patch_size, stride, dilation):
        patches = []

        for i in range(0, input_image.shape[2] - (patch_size[0] - 1), stride):
            for j in range(0, input_image.shape[3] - (patch_size[1] - 1), stride):
                x_range = (patch_size[0] + ((patch_size[0] - 1) * (dilation - 1)))
                y_range = (patch_size[1] + ((patch_size[1] - 1) * (dilation - 1)))
                #print(x_range, y_range)
                if i + x_range <= input_image.shape[2] and j + y_range <= input_image.shape[3]:
                    coord = torch.Tensor([x * input_image.shape[2] + y for x in range(i, i + x_range, dilation) for y in
                                      range(j, j + y_range, dilation)])


                patches.append(coord)

        indexes_of_patches = torch.stack(patches, dim=0).long()
        return indexes_of_patches

    def image_to_vectors(self, input_image, indexes_of_patches):
        p, num_pixels = indexes_of_patches.shape
        indexes_of_patches = indexes_of_patches.view(p * num_pixels)



        out = input_image.view(input_image.shape[0], input_image.shape[1], -1)
        #print(out.shape)
        out = torch.index_select(out, dim=2,
                                 index=indexes_of_patches.to(input_image.device))
        out = out.view(out.shape[0], out.shape[1], p, num_pixels)

        return out

    def build_block(self):
        x = torch.zeros(self.input_shape)

        if type(self.kernel_size) is int:
            self.kernel_size = (self.kernel_size, self.kernel_size)

        out = x.view(-1, x.shape[2], x.shape[3], x.shape[4])


        if self.padding > 0:
            out = F.pad(out, pad=[self.padding, self.padding, self.padding, self.padding], mode='constant', value=0)

        self.indexes_of_patches = self.image_to_patch_coordinates(input_image=out, patch_size=self.kernel_size, stride=self.stride, dilation=self.dilation).to(x.device)

        x_vectors = self.image_to_vectors(out, indexes_of_patches=self.indexes_of_patches)
        self.spatial_shape = None
        print('block built', x_vectors.shape)

    def forward(self, x, weights, biases=None):

        # assert x.shape[0] == weights.shape[0], \
        #     "The batch size of the input images needs to be equal to the batch size of the params"
        # if biases is not None:
        #     assert weights.shape[0] == biases.shape[0], \
        #     "The batch size of the weight parameters needs to be equal to the batch size of the bias parameters"
        weights = weights.permute([0, 2, 3, 4, 1]).view(weights.shape[0], -1, weights.shape[1])

        b_out, b_in = x.shape[:2]
        out = x.view(-1, x.shape[2], x.shape[3], x.shape[4])

        if self.padding > 0:
            out = F.pad(out, pad=[self.padding, self.padding, self.padding, self.padding], mode='constant', value=0)


        out = self.image_to_vectors(out, indexes_of_patches=self.indexes_of_patches)

        out = torch.cat(out.unbind(1), 2)
        out = out.view(b_out, b_in * out.shape[1], out.shape[2])

        out = torch.bmm(out, weights)


        if biases is not None:
            out = out + biases.unsequeeze(1)

        out = out.view(b_out, b_in, -1, out.shape[-1])

        if self.spatial_shape is None:
            self.spatial_shape = int(np.floor(np.sqrt(out.shape[2])))

        out = out.view(b_out, b_in, self.spatial_shape, self.spatial_shape, out.shape[-1])

        out = out.permute([0, 1, 4, 2, 3])

        return out

[ssnl]

Can't this be done efficiently using group conv?

[AntreasAntoniou]

Can you elaborate? With standard convolutions, simply appending more filter parameters wouldn't work because the whole batch given would be convolved with all filters. So in detail:
Let's assume that I have:

A collection of batches of images of shape (b_i, b_j, c, h, w)
A batch of kernel weights of shape (b_i, out_features, in_features, kernel_height, kernel_width)
A stride of 1, kernel size of 3 and padding of 0.

If I were to simply create a batch of size (b_i * b_j, c, h, w), and weight kernels of size (b_i * out_features, in_features, kernel_height, kernel_width)

Due to the kernel size, each dot product generating a feature would be of size (9x1). And, we would end up computing (b_i * b_j * b_i * f_out) * ((h - 2) * (w - 2)) dot products. From that, one would have to throw away the features generated from convolving batches with kernel parameters they shouldn't have been convolved with. This is approximately ((b_i * b_j * (b_i - 1) * f_out) * ((h - 2) * (w - 2)) unnecessary dot products. As b_i increases so do the wasted dot products.

Now, as I understand them, group convolutions, can convolve the same batch by different filters, depending on the num_groups and the num_input_channels.

Can you use the above framework to explain your proposal?

[ssnl]

with slight reordering,

(b_j, b_i, c_in, h, w) -> (b_j, b_i * c_in, h, w)
(b_i, c_out, c_in, kh, kw) -> (b_i * c_out, c_in, kh, kw)

then do group=b_i conv.

[AntreasAntoniou]

@ssnl You are correct. This is a very efficient way of doing this. I wasn't very aware of the details behind pytorch's groups argument. I knew what group convs were, mathematically and intuitively. However, I've never actually made use of them in Pytorch. Thanks for your help with this. I will implement BatchConv2D and BatchConv1D using group convs, and then BatchLinearLayer using torch.bmm. I will post the code here once I've stress tested them with some run-speed measurements.

P.S. I re-derived the reordering/reshape (because I like to understand what I am doing). I managed to arrive at the same reordering idea once I understood the group conv argument. However, I was wondering where your 2 * b_i came from, since I had b_i. I am glad to see that you edited your answer and we are actually in agreement. :)

[ssnl]

@AntreasAntoniou Awesome. Yeah the 2* factor was a typo. Glad that you figured it out. :)

[AntreasAntoniou]

import torch
import torch.nn as nn
import torch.nn.functional as F


class BatchConv2DLayer(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1,
                 padding=0, dilation=1):
        super(BatchConv2DLayer, self).__init__()

        self.stride = stride
        self.padding = padding
        self.dilation = dilation

        self.in_channels = in_channels
        self.out_channels = out_channels

    def forward(self, x, weight, bias=None):
        if bias is None:
            assert x.shape[0] == weight.shape[0], "dim=0 of x must be equal in size to dim=0 of weight"
        else:
            assert x.shape[0] == weight.shape[0] and bias.shape[0] == weight.shape[
                0], "dim=0 of bias must be equal in size to dim=0 of weight"

        b_i, b_j, c, h, w = x.shape
        b_i, out_channels, in_channels, kernel_height_size, kernel_width_size = weight.shape

        out = x.permute([1, 0, 2, 3, 4]).contiguous().view(b_j, b_i * c, h, w)
        weight = weight.view(b_i * out_channels, in_channels, kernel_height_size, kernel_width_size)

        out = F.conv2d(out, weight=weight, bias=None, stride=self.stride, dilation=self.dilation, groups=b_i,
                       padding=self.padding)


        out = out.view(b_j, b_i, out_channels, out.shape[-2], out.shape[-1])

        out = out.permute([1, 0, 2, 3, 4])

        if bias is not None:
            out = out + bias.unsqueeze(1).unsqueeze(3).unsqueeze(3)

        return out

class BatchConv1DLayer(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1,
                 padding=0, dilation=1):
        super(BatchConv1DLayer, self).__init__()

        self.stride = stride
        self.padding = padding
        self.dilation = dilation

        self.in_channels = in_channels
        self.out_channels = out_channels

    def forward(self, x, weight, bias=None):
        if bias is None:
            assert x.shape[0] == weight.shape[0], "dim=0 of x must be equal in size to dim=0 of weight"
        else:
            assert x.shape[0] == weight.shape[0] and bias.shape[0] == weight.shape[
                0], "dim=0 of bias must be equal in size to dim=0 of weight"

        b_i, b_j, c, h = x.shape
        b_i, out_channels, in_channels, kernel_width_size = weight.shape

        out = x.permute([1, 0, 2, 3]).contiguous().view(b_j, b_i * c, h)
        weight = weight.view(b_i * out_channels, in_channels, kernel_width_size)

        out = F.conv1d(out, weight=weight, bias=None, stride=self.stride, dilation=self.dilation, groups=b_i,
                       padding=self.padding)

        out = out.view(b_j, b_i, out_channels, out.shape[-1])

        out = out.permute([1, 0, 2, 3])

        if bias is not None:
            out = out + bias.unsqueeze(1).unsqueeze(3)

        return out

class BatchLinearLayer(nn.Module):
    def __init__(self):
        super(BatchLinearLayer, self).__init__()

    def forward(self, x, weight, bias=None):
        if bias is None:
            assert x.shape[0] == weight.shape[0], "dim=0 of x must be equal in size to dim=0 of weight"
        else:
            assert x.shape[0] == weight.shape[0] and bias.shape[0] == weight.shape[
                0], "dim=0 of bias must be equal in size to dim=0 of weight"

        out = torch.bmm(x, weight)

        if bias is not None:
            out = out + bias.unsqueeze(1)

        return out

[AntreasAntoniou]

I've just finished implementing the batch operator layers, along with a short usage script found below:

import torch
import torch.nn as nn

from batch_operators import BatchConv2DLayer, BatchConv1DLayer, BatchLinearLayer

x = torch.ones(100, 50, 3, 32, 32)
params = torch.empty(100, 64, 3, 3, 3)
bias = torch.zeros(100, 64)
nn.init.xavier_uniform_(params)

layer = BatchConv2DLayer(in_channels=3, out_channels=32, stride=1, padding=0, dilation=1)

out = layer.forward(x=x, weight=params, bias=bias)

print(out.shape)

x = torch.ones(100, 50, 3, 32)
params = torch.empty(100, 64, 3, 3)
bias = torch.zeros(100, 64)
nn.init.xavier_uniform_(params)

layer = BatchConv1DLayer(in_channels=3, out_channels=32, stride=1, padding=0, dilation=1)

out = layer.forward(x=x, weight=params, bias=bias)

print(out.shape)

x = torch.ones(100, 10, 50)
params = torch.empty(100, 50, 75)
bias = torch.zeros(100, 75)
nn.init.xavier_uniform_(params)

layer = BatchLinearLayer()

out = layer.forward(x=x, weight=params, bias=bias)

print(out.shape)

[AntreasAntoniou]

I'll now measure the speed of the parallel variant vs the sequential variant, on both high and low b_i.

[jramapuram]

@AntreasAntoniou @ssnl : thanks for solving this.

For the lazy folk like myself that don't want to pass around weights I have the following:

class BatchConv2D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size,
                 stride=1, padding=0, dilation=1, groups=1, bias=True):
        """ Batch convolve a set of inputs by groups in parallel. Similar to bmm.

        :param in_channels: (b_j, b_i, c_in, h, w) where b_j are the parallel convs to run
        :param out_channels: output channels from conv
        :param kernel_size: size of conv kernel
        :param stride: the stride of the filter
        :param padding: the padding around the input
        :param dilation: the filter dilation
        :param groups: number of parallel ops
        :param bias: whether of not to include a bias term in the conv (bool)
        :returns: tensor of (b_i, c_out, c_in, kh, kw) with batch convolve done
        :rtype: torch.Tensor

        """
        super(BatchConv2D, self).__init__()
        self.out_channels = out_channels
        self.conv = nn.Conv2d(in_channels*groups, out_channels*groups,
                              kernel_size, stride=stride,
                              padding=padding, dilation=dilation,
                              groups=groups, bias=bias)

    def forward(self, x):
        """ (b_j, b_i, c_in, h, w) -> (b_j, b_i * c_in, h, w) --> (b_j, b_i, c_out, h, w)

        :param x: accepts an input of (b_j, b_i, c_in, h, w) where b_j are the parallel groups
        :returns: (b_j, b_i , c_out, kh, kh, kw)
        :rtype: torch.Tensor

        """
        assert len(x.shape) == 5, "batchconv2d expects a 5d [{}] tensor".format(x.shape)
        b_i, b_j, c, h, w = x.shape
        out = self.conv(x.permute([1, 0, 2, 3, 4]).contiguous().view(b_j, b_i * c, h, w))
        return out.view(b_j, b_i, self.out_channels,
                        out.shape[-2], out.shape[-1]).permute([1, 0, 2, 3, 4])

Example usage:

model = BatchConv2D(3, 32, 4, groups=2)
print(model(torch.randn(2, 5, 3, 32, 32)).shape)

Result:

torch.Size([2, 5, 32, 29, 29])

Edit: I believe nn.Linear doesn't need anything to work out-of-the-box:

nn.Linear(3, 32)(torch.randn(2, 5, 3)).shape

Output:

torch.Size([2, 5, 32])

This is because the underlying F.linear (and corresponding torch.matmul) accepts batch dimensions:

        - Input: :math:`(N, *, in\_features)` where `*` means any number of
          additional dimensions

[li-roy]

This looks addressed, so I'm going to close it. Feel free to reopen if I'm mistaken.

[AntreasAntoniou]

@li-roy I guess the question here is whether something like this should be part of the pytorch Functional module.