CODEWITHSUNDEEP
pytorch
BlackSith
BlackSith

Reputation: 87

Get the mean of each column for a 2D PyTorch tensor object

I want to know what is the most simple way to get the mean of the matrix along each column, namely my tensor only has two dimensions with shape (m X n).

For example, if I have a tensor object

T = torch.FloatTensor([[2.6, 5], [1.7, 6], [3.2, 7], [2.1, 8]])

I want some function that could return a tensor object as following

([2.4, 6.5])

I know I could do with for loop, but I wonder if there is a PyTorch built-in method that could do this for me, or is there some more elegant way for this purpose.

for i in range(T.size(1)):
  mean[i] = torch.mean(T[:, i])

Thanks for any help.

Upvotes: 3

Views: 6060

Answers (1)

Shai
Shai

Reputation: 114826

Please read the doc on torch.mean. It has an optional parameter dim:

dim (int or tuple of ints) – the dimension or dimensions to reduce.

In your case:

T.mean(dim=0)

Will give you the mean along the first dim.

Note that many other "reduction" operations in pytorch (and in ) also have this option to reduce along a specific dimension only (e.g., std, sum etc.).

Upvotes: 2

Related Questions