Hello all,
I have a network that generates kernel weights for a 2D convolution operation. It takes a single input and generates a weight vector which then reshaped into KxK kernel where K is the kernel size. Finally, I apply this kernel to an image. I’m trying to make this work with batches. So when the network generates a [N, KxK] weight vector, I would like to have a tensor with size [N, K, K] that is applied to [N, 1, H, W] images. Here, N is the batch size, H and W are the image sizes. I can do this with for loop over the batches but I’m wondering if there is a more efficient way. Below is my network:
class FilterNet(nn.Module):

    def __init__(self, obs_size, hidden_size, kernel_size, activation_func='tanh'):
        super(FilterNet, self).__init__()
        self.kernel_size = kernel_size

        self.fc_obs1 = nn.Linear(obs_size, hidden_size)
        self.fc_obs2 = nn.Linear(hidden_size, hidden_size)
        self.fc_obs3 = nn.Linear(hidden_size, hidden_size)
        self.fc_obs4 = nn.Linear(hidden_size, kernel_size * kernel_size)

        self.act_func = getattr(torch, activation_func)

    def forward(self, images, tau):
        hidden = self.act_func(self.fc_obs1(tau))
        hidden = self.act_func(self.fc_obs2(hidden))
        hidden = self.act_func(self.fc_obs3(hidden))
        kernel_weigths = self.fc_obs4(hidden)

        heatmaps = []
        for weights_i, image_i in zip(kernel_weigths.chunk(kernel_weigths.size(0), dim=0), images.chunk(images.size(0), dim=0)):
            kernel = weights_i.view(1, 1, self.kernel_size, self.kernel_size)
            heatmap = F.conv2d(image_i, kernel, padding=2)
            heatmap = heatmap.squeeze()
            heatmaps += [heatmap]
        heatmaps = torch.stack(heatmaps, 0)

        # kernel_weigths = kernel_weigths.view(kernel_weigths.size(0), 1, self.kernel_size, self.kernel_size)
        # heatmaps = F.conv2d(images, kernel_weigths, padding=2)
        
        return heatmaps

Note the commented section of the code does not work.
Thanks!