CODEWITHSUNDEEP
pythonpytorch
stacker_bell
stacker_bell

Reputation: 1

2-D Tensor calculated by the mean of 3-D Tensor by specific dimension

I have a 3-D tensor, with shape (3000, 20, 5). I want to create a 2-D tensor, of shape (3000, 5), using the mean values of the second dimension of the 3-D tensor.

So basically, I want to perform something like:

mean_value = torch.mean(3d_tensor[0][:][0])

But getting values for all values of dimension one a three. I could do a for loop, e.g.:

for j in range(0, 3d_tensor.size()[2]):
    for i in range(0, len(3d_tensor)):

        mean_values[j][i] = torch.mean(3d_tensor[i][:][j])

But this takes a long time to process for large amounts of data.

Upvotes: 0

Views: 217

Answers (1)

Szymon Maszke
Szymon Maszke

Reputation: 24726

You could simply specify axis along which mean should be taken:

mean = torch.mean(tensor, dim=1)

This gives you data of shape (3000, 5)

Upvotes: 1

Related Questions