I want to achieve a neural network where, on a given layer, the layer operation applied is not the same on all instances of a batch. For instance, if my network is a MLP, I want each instance’s feature vector to be multiplied by a different matrix. In that case, I can get that result by using torch.matmul and extending the weight tensor across the batch dimension.
When I try to do the same with a CNN, I can’t. nn.functionnal.conv2d does not accept a weight tensor that has an extra dimension for batch. To get around this, I compute each convolution sequentially and stack them. While it works, it does not use my GPU efficiently and is thus really slow.
My question regarding this is the following: is there something in the pytorch API that I missed and that I could use to achieve instance-dependent convolution, similar to how matmul functions? If not, is there any particular reason why?