How to apply a nn.Module (i.e. CNN) across an axis (i.e. Video input) in a parallelizable way

[seunggs]

Hi, I’m trying to apply CNN to each image in a video. Currently, my implementation uses a for loop and torch.cat where I take each image and apply the CNN module in the loop. But clearly, this is sequential and I don’t see why it can’t be parallelized in theory since all images are independent from each other.

However, I’m not sure how this can be accomplished. I couldn’t find any built-in function for PyTorch. Is there a way to do this in parallel in PyTorch Lightning?

My video input shape looks like this: (batch_size, seq_len, channel, height, width) and CNN takes input shape of (batch_size, channel, height, width).

Thanks in advance for your help!

[SkafteNicki]

You can simply convert your (batch_size, seq_len, channel, height, width) tensor into an (batch_size*seq_len, channel, height, width) tensor, run your model and then reshape your output back:

batch_size, seq_len, channel, height, width = 5, 10, 3, 28, 28 # just random picked
input = torch.randn(batch_size, seq_len, channel, height, width)
input = input.reshape(batch_size * seq_len, channel, height, width)
output = model(input) 
# split the batch dimension back into the original batch size and sequence length
output = output.reshape(batch_size, seq_len, *output.shape[1:])

[seunggs]

Oh right - this should work. Thanks a lot for taking a look at this @SkafteNicki!