How to compute stride mean in Pytorch?

Question

For example, we have x = [1, 2, 3, 4, 5] with stride = 3, step = 1, then the stride mean would be result = [2, 3, 4]. More specifically, I have a vector with shape = (batch_size, sequence_len, embed_dim), and I want to perform the above operation on dim = 1 (sequence_len). How can I implement that efficiently in PyTorch?

Answer 1

I think what you mean - in PyTorch notation - is a kernel size of 3 and a stride of 1. You can use torch.nn.AvgPool1d to perform this kind of operation:

mean = nn.AvgPool1d(kernel_size=3, stride=1)

Note, you will need one extra dimension, for the channel, to be compatible with this kind of layer:

>>> x = torch.tensor([[1, 2, 3, 4, 5]]).unsqueeze(0)
tensor([[[1, 2, 3, 4, 5]]])

>>> mean(x)
tensor([[[2, 3, 4]]])