How to compute the mean and standard deviation of a tensor in PyTorch?

A PyTorch tensor is similar to a NumPy array but optimized for GPU acceleration. Computing mean and standard deviation are fundamental statistical operations in deep learning for data normalization and analysis.

Basic Syntax

PyTorch provides built-in functions for these statistical computations ?

• torch.mean(input, dim=None) ? Computes the mean value
• torch.std(input, dim=None) ? Computes the standard deviation

Computing Mean and Standard Deviation of 1D Tensor

Let's start with a simple one-dimensional tensor ?

import torch

# Create a 1D tensor
tensor_1d = torch.tensor([2.453, 4.432, 0.754, -6.554])
print("Tensor:", tensor_1d)

# Compute mean and standard deviation
mean_val = torch.mean(tensor_1d)
std_val = torch.std(tensor_1d)

print("Mean:", mean_val)
print("Standard deviation:", std_val)

Tensor: tensor([ 2.4530,  4.4320,  0.7540, -6.5540])
Mean: tensor(0.2713)
Standard deviation: tensor(4.7920)

Computing Statistics for 2D Tensor

For multi-dimensional tensors, you can compute statistics across all elements or along specific dimensions ?

import torch

# Create a 3x4 2D tensor
tensor_2d = torch.tensor([[2., 4., 7., -6.],
                          [7., 33., -62., 23.],
                          [2., -6., -77., 54.]])
print("Tensor:")
print(tensor_2d)

# Overall mean and standard deviation
overall_mean = torch.mean(tensor_2d)
overall_std = torch.std(tensor_2d)
print("\nOverall Mean:", overall_mean)
print("Overall Standard deviation:", overall_std)

Tensor:
tensor([[  2.,   4.,   7.,  -6.],
        [  7.,  33., -62.,  23.],
        [  2.,  -6., -77.,  54.]])

Overall Mean: tensor(-1.5833)
Overall Standard deviation: tensor(36.2703)

Dimension-wise Computations

Use the dim parameter to compute statistics along specific dimensions ?

import torch

tensor_2d = torch.tensor([[2., 4., 7., -6.],
                          [7., 33., -62., 23.],
                          [2., -6., -77., 54.]])

# Column-wise statistics (dim=0)
col_mean = torch.mean(tensor_2d, dim=0)
col_std = torch.std(tensor_2d, dim=0)
print("Column-wise Mean:", col_mean)
print("Column-wise Std:", col_std)

# Row-wise statistics (dim=1)
row_mean = torch.mean(tensor_2d, dim=1)
row_std = torch.std(tensor_2d, dim=1)
print("\nRow-wise Mean:", row_mean)
print("Row-wise Std:", row_std)

Column-wise Mean: tensor([  3.6667,  10.3333, -44.0000,  23.6667])
Column-wise Std: tensor([ 2.8868, 20.2567, 44.7996, 30.0056])

Row-wise Mean: tensor([ 1.7500,  0.2500, -6.7500])
Row-wise Std: tensor([ 5.5603, 42.8593, 53.8602])

Key Parameters

Conclusion

Use torch.mean() and torch.std() to compute tensor statistics. Specify dim parameter for dimension-wise calculations, which is essential for batch processing in neural networks.