Python

PyTorch – Median()

We will see how to return the median of a tensor using median() in this PyTorch tutorial.

PyTorch is an open-source framework available with a Python programming language. Tensor is a multidimensional array that is used to store the data. To use a tensor, we have to import the torch module. To create a tensor, the method used is tensor().

Syntax:

torch.tensor(data)

Where the data is a multi-dimensional array.

Median()

Median() in PyTorch is used to return the median of the elements present in the input tensor object.

Syntax:

torch.median(tensor,dim)

Where:

1. The tensor is the input tensor.

2. Dim is to reduce the dimension. Dim=0 specifies the column comparison which gets the median along a column and dim=1 specifies the row comparison which gets the median along the row.

Return:

It returns the median along with the index position where it is present in the tensor.

Example 1:

In this example, we will create a tensor with 2 dimensions that has 3 rows and 5 columns and apply the median() function on rows and columns.

#import torch module

import torch

 

#create a tensor with 2 dimensions (3 * 5)

#with random elements using randn() function

data = torch.randn(3,5)

 

#display

print(data)

 

print()

 

#get Median along columns

print("Median across columns:")

print(torch.median(data, dim=0))

 

print()

 

#get Median along rows

print("Median across rows:")

print(torch.median(data, dim=1))

Output:

tensor([[ 0.9562, 0.4517, 2.1758, -0.7846, -0.7103],

[-0.4080, 1.9359, 1.0324, 0.0120, -0.4021],

[ 0.6448, -0.6840, 1.5963, 0.4659, 0.5414]])

Median across columns:

torch.return_types.median(

values=tensor([ 0.6448, 0.4517, 1.5963, 0.0120, -0.4021]),

indices=tensor([2, 0, 2, 1, 1]))

Median across rows:

torch.return_types.median(

values=tensor([0.4517, 0.0120, 0.5414]),

indices=tensor([1, 3, 4]))

We can see that the median is returned across the columns and rows along with the index positions (indices).

Example 2:

Create a Tensor with 5 * 5 matrix and return the median across the rows and columns.

#import torch module

import torch

 

#create a tensor with 2 dimensions (5 * 5)

#with random elements using randn() function

data = torch.randn(5,5)

 

#display

print(data)

 

print()

 

#get Median along columns

print("Median across columns:")

print(torch.median(data, dim=0))

 

print()

 

#get Median along rows

print("Median across rows:")

print(torch.median(data, dim=1))

Output:

tensor([[-1.0895, 1.2629, 0.5036, -0.8317, 0.9415],

[-0.1067, 0.5748, 0.1773, 0.1904, 0.1359],

[-0.7851, -0.1791, -0.1605, 1.8524, 0.3375],

[-0.6451, 0.5093, -2.0334, -0.6462, -0.5074],

[-0.7813, -1.5316, 0.3567, 1.5330, 0.6483]])

Median across columns:

torch.return_types.median(

values=tensor([-0.7813, 0.5093, 0.1773, 0.1904, 0.3375]),

indices=tensor([4, 3, 1, 1, 2]))

Median across rows:

torch.return_types.median(

values=tensor([ 0.5036, 0.1773, -0.1605, -0.6451, 0.3567]),

indices=tensor([2, 2, 2, 0, 2]))

We can see that the median across the rows and columns were returned along with the indices.

Without the Dim Parameter

If we don’t specify the dim parameter, it returns the whole median.

Example 1:

Create a 2D tensor with 5*5 matrix and return the median.

#import torch module

import torch

#create a tensor with 2 dimensions (5 * 5)

#with random elements using randn() function

data = torch.randn(5,5)

#display

print(data)

print()

#get median

print("Median :")

print(torch.median(data))

Output:

tensor([[ 0.7692, 1.9155, 1.0876, 0.6794, 0.5791],

[-0.2212, -1.4801, 0.5487, -0.5765, -0.4750],

[-0.0681, 0.5216, 1.1400, -0.1743, 0.0344],

[-0.4639, -1.2341, -1.0520, 0.1718, -0.1909],

[ 0.6911, -1.4869, 2.7762, -0.3645, -0.4775]])

Median :

tensor(-0.0681)

Example 2:

Create a 1D tensor with 5 values and return the median.

#import torch module

import torch

 

#create a tensor with 5 numeric values

data = torch.tensor([10.6,20.7,30.6,40.4,50.0])

 

#display

print(data)

 

print()

 

#get median

print("Median :")

print(torch.median(data))

Output:

tensor([10.6000, 20.7000, 30.6000, 40.4000, 50.0000])

Median :

tensor(30.6000)

Work with CPU

If you want to run a median() function on the CPU, we have to create a tensor with a cpu() function. This will run on a CPU machine.

When we create a tensor, this time, we can use the cpu() function.

Syntax:

torch.tensor(data).cpu()

Example 1:

In this example, we will create a tensor with 2 dimensions that has 3 rows and 5 columns with the cpu() function and apply the median() on rows and columns.

#import torch module

import torch

 

#create a tensor with 2 dimensions (3 * 5)

#with random elements using randn() function

data = torch.randn(3,5).cpu()

 

#display

print(data)

 

print()

 

#get Median along columns

print("Median across columns:")

print(torch.median(data, dim=0))

 

print()

 

#get Median along rows

print("Median across rows:")

print(torch.median(data, dim=1))

Output:

tensor([[ 0.9872, 0.1258, -0.0952, 0.3269, -1.6033],

[-0.2432, -1.0049, -0.9058, 0.9438, 0.3060],

[-2.8293, 1.4515, -0.9482, 0.9876, 0.2767]])

Median across columns:

torch.return_types.median(

values=tensor([-0.2432, 0.1258, -0.9058, 0.9438, 0.2767]),

indices=tensor([1, 0, 1, 1, 2]))

Median across rows:

torch.return_types.median(

values=tensor([ 0.1258, -0.2432, 0.2767]),

indices=tensor([1, 0, 4]))

We can see that the median is returned across the columns and rows with indices.

Example 2:

Create a Tensor with 5 * 5 matrix with the cpu() function and return the median across the rows and columns.

#import torch module

import torch

 

#create a tensor with 2 dimensions (5 * 5)

#with random elements using randn() function

data = torch.randn(5,5).cpu()

 

#display

print(data)

 

print()

 

#get Median along columns

print("Median across columns:")

print(torch.median(data, dim=0))

 

print()

 

#get Median along rows

print("Median across rows:")

print(torch.median(data, dim=1))

Output:

tensor([[-0.3739, -1.2500, -1.9125, -0.4597, 0.2058],

[-0.1885, -0.4993, -1.0801, -0.1367, -0.5683],

[-0.1242, 0.1221, -0.2267, -0.7851, 0.6797],

[ 2.2487, 0.0141, 0.1632, -0.4924, -0.9134],

[-1.6101, 0.5051, -0.2004, -0.4901, -0.3358]])

Median across columns:

torch.return_types.median(

values=tensor([-0.1885, 0.0141, -0.2267, -0.4901, -0.3358]),

indices=tensor([1, 3, 2, 4, 4]))

Median across rows:

torch.return_types.median(

values=tensor([-0.4597, -0.4993, -0.1242, 0.0141, -0.3358]),

indices=tensor([3, 1, 0, 1, 4]))

We can see that the median across the rows and columns were returned.

Conclusion

In this PyTorch lesson, we learned about the median() function and how to apply it on a tensor to return a median across the columns and rows.

We also created a tensor with the cpu() function and returned the median. If the dim is not specified in two or multi-dimensional tensor, it returns the whole median.

About the author

Gottumukkala Sravan Kumar

B tech-hon's in Information Technology; Known programming languages - Python, R , PHP MySQL; Published 500+ articles on computer science domain