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.