Ecanyte
Ecanyte

Reputation: 81

Combine features in tensor in pytorch

I have a tensor with dim 2n x m. I want to compute an output tensor with dim n x m, where the i-th and the i+1-th entry are added together and divided by 2, i.e. (f_0, f_1, f_2, f_3, ...) -> ((f_0+f_1)/2, (f_2+f_3)/2, ...). How can I achieve this without looping over the tensor?

Thanks for your help.

Upvotes: 1

Views: 253

Answers (1)

joe32140
joe32140

Reputation: 1392

I would reshape the tensor to (n,2,m) and take the mean of dim 1.

In [7]: x = torch.arange(12).view(4,3).float()

In [8]: x
Out[8]: 
tensor([[ 0.,  1.,  2.],
        [ 3.,  4.,  5.],
        [ 6.,  7.,  8.],
        [ 9., 10., 11.]])

In [9]: x.view(2,2,3).mean(dim=1)
Out[9]: 
tensor([[1.5000, 2.5000, 3.5000],
        [7.5000, 8.5000, 9.5000]])

Upvotes: 1

Related Questions