Hi! Now I want to pass a 5-d tensor to conv layers, where the the tensor is of [k, batch_size, C,H,W]. The later 4 dimension [batch_size,C,H,W] is just the usual 4-D tensor of image input, but this 4-D tensor is copied and expanded k times along the first dimension. Such an input is necessary in models like IWAE. I want to do usual 2D convolution on this tensor, while making sure that the forward pass output for each of the k copies are the same. I think reshaping it to [k*batch_size,C,H,W] will make the forward output for each copy different, so is there any smart way to implement this?
I think reshaping it to k*batch_size,C,H,W would work. you can check this out with for example batch_size =1, and k=3 , the output would be 3,C,H,W then, and you can verify that 0,:,:,: is the same as 1,:,:,:
1 Like