Pytorch reshape to modify batch size in forward()

Question

I have a tensor of shape, say: [4,10] where 4 is the batch size and 10 is the length of my input samples buffer. Now, I know that it is really [4,5+5] i.e. the input samples buffer consists of two windows of length 5 which can be processed independently and, best, in parallel. What I am doing is, inside forward() of my model I first reshape the tensor to [8,5], run my layers on it, and then reshape it back to [4,-1] and return. What I am hoping to get from this is Pytorch would run my model on each of the windows (kind of sub-batches) in parallel, effectively yielding a parallel-for loop. It runs OK, Pytorch does not complain or anything but I am getting weird results. I'd like to know if Pytorch can work this way before I dive into debugging my model.

Yes there's nothing absurd in what you described, pytorch should be able to work this way. If you provide some code samples with inputs and expected outputs we could help you debug all this — trialNerror

cpajdel cpajdel · Accepted Answer · 2021-01-12T09:00:11

Well, it doesn't. The reason being the ordering pytorch uses for tensor reshaping. This can be seen by running the small repro code below. It would be nice to have something like a 'rebatch' function in pytorch that would take care of proper memory layout as a foundation for parallel-for construcs (provided it can even be done memory-efficiently in generic case).

import torch

conv = torch.nn.Conv1d(1,3,1)

def conv_batch(t, conv, window):
    batch = t.shape[0]
    t = t.view(-1, t.shape[1], window)
    t = conv(t)
    t = t.view(batch, t.shape[1], -1)
    return t

batch = 1
channels = 1
width = 4

window = 2

x = torch.arange(batch*channels*width)
x = x.view(batch,channels,width).float()

r1 = conv(x)
r2 = conv_batch(x, conv, window)

print(r1)
print(r2)
print(r1==r2)

Pytorch reshape to modify batch size in forward()

1 Answers