如何计算自定义数据集的统计分布（均值、标准差）及代码验证

原代码的问题 #

你的代码存在两处核心逻辑错误，导致计算结果不准确：

1. 均值计算逻辑错误：images.mean(0).sum(0)的操作偏离了正确统计逻辑——对于形状为(batch, channel, H, W)的图像张量，需要累加每个通道的所有像素总和，而非先对batch维度求均值再对通道维度求和。
2. 标准差计算逻辑错误：直接对每个batch的标准差累加再除以样本数是错误的，标准差是基于均值的二次统计量，不能通过简单平均batch的标准差得到整体结果。

正确的实现代码 #

import torch
from torch.utils.data import DataLoader

# 假设已有自定义数据集对象dataset
batch_size = 64
# 计算统计量无需shuffle，关闭可节省资源
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)

# 根据图像通道数初始化，3通道用torch.zeros(3)，单通道用torch.zeros(1)
mean = torch.zeros(3)
std = torch.zeros(3)
total_pixels = 0  # 总像素数 = 样本数 × 图像高度 × 图像宽度

for images, _ in data_loader:
    # 获取当前batch的张量维度
    batch_size = images.size(0)
    h, w = images.size(2), images.size(3)
    
    # 累加每个通道的像素总和（对batch、H、W维度求和）
    channel_sum = images.sum([0, 2, 3])
    mean += channel_sum
    
    # 累加每个通道的像素平方总和
    channel_sum_sq = (images ** 2).sum([0, 2, 3])
    std += channel_sum_sq
    
    total_pixels += batch_size * h * w

# 计算全局均值
mean /= total_pixels
# 利用方差公式计算全局标准差：sqrt(E[X²] - (E[X])²)
std = torch.sqrt(std / total_pixels - mean ** 2)

print("Mean:", mean)
print("Standard Deviation:", std)

关键说明 #

• 均值计算：通过累加每个通道的所有像素总和，再除以数据集的总像素数，得到每个通道的真实均值。
• 标准差计算：借助方差的数学公式Var(X) = E[X²] - (E[X])²，先计算像素平方的期望，减去均值的平方得到方差，最后开根号得到标准差，这是计算总体标准差的标准方法。
• 注意事项：如果数据集图像尺寸不统一，需确保预处理时将图像resize到相同尺寸，否则无法正确统计；单通道图像只需将初始化的torch.zeros(3)改为torch.zeros(1)即可。

内容的提问来源于stack exchange，提问作者linkho