利用 torch.mean()计算图像数据集的均值和标准差

本文参考了youtube博主Aladdin Persson和Kaggle用户Afia Ibnath的方法

torch.mean()

Pytorch官方文档

Pytorch官方文档

参数：

input: 输入，接收 tensor 类型
dim: 维度

Example:

输出结果：

*
由结果可以看出
dim默认为计算全部均值
当dim = 1 时跨列求平均
当dim = 0是跨行求平均
*

计算方差：

def get_mean_std(loader): 
    # var[X] = E[X**2] - E[X]**2 方差公式， var[]代表方差，E[]表示期望(平均值)
    channels_sum, channels_sqrd_sum, num_batches = 0, 0, 0
    for data, _ in tqdm(loader):
        channels_sum += torch.mean(data, dim = [1, 2]) 
        channels_sqrd_sum += torch.mean(data ** 2, dim = [1, 2])
        num_batches += 1


    mean = channels_sum / num_batches
    std = (channels_sqrd_sum / num_batches - mean ** 2) ** 0.5

    return mean, std, num_batches

该函数可以用来计算 RGB 三通道数据集的均值和方差。
即插即用，非常方便。

运行结果

此为计算猫狗训练集的均值和方差
训练数据集：
https://www.kaggle.com/tongpython/cat-and-dog/code

注意

1.函数接收的图片为tensor类型，需要提前转换图片类型为（channel，H，W）
2.如果传入图片形式为（batchsize，channel，H，W）需要把参数dim改变为[0， 2，3]。

这是本人第一次发表文章，如有不对之处还望大家批评斧正。未经允许，谢绝转载！