I’m implementing a custom forward function for e.g. nn.Conv2d and nn.Linear but I want to do 1 forward pass with N slightly noised versions of the weight for each sample in the mini-batch, without having to break it down into N different forward passes. For example, I might do:
for p in range(perturbations):
features.append(F.conv2d(input[p, :], weight+noise[p], bias, stride, padding, dilation, groups))
So input would be of dimensions B x C_in x H x W whereas the noised-weights would be of dimensions N x C_out x C_in x k x k where N = B in the above example. How can I vectorise the forward passes? Thanks!
Expand your input to create N copies of itself, packaged in the in_channels
dimension, so that it has dimensions B x N * C_in x H x W. Then use groups = N in your call to conv2d() and package your N noisy copies of your weight tensor into a single weight tensor with dimensions, similar to as you
described, of N * C_out x C_in x k x k. You will also have to expand()
the bias you pass in to replicate it N-fold to dimensions N * C_out.