计算图像数据集三通道的均值和方差

在深度学习与图像有关的任务中，我们经常需要计算图像数据集（特别是私有数据集）三通道的均值和方差来对图像进行norm处理。下面以CIFAR10数据集为例来展示如何计算图像数据集三通道的均值和方差：

import torch
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
import torchvision.transforms as transforms

def get_mean_std(loader):
    channels_sum, channels_squared_sum, num_batches = 0, 0, 0
    for data, _ in loader:
        channels_sum += torch.mean(data, dim=[0, 2, 3])
        channels_squared_sum += torch.mean(data ** 2, dim=[0, 2, 3])
        num_batches += 1
    mean = channels_sum / num_batches
    mean_squared = channels_squared_sum / num_batches
    #VAR[X]=E[X**2]-E(X)**2
    std = (mean_squared - mean ** 2) ** 0.5
    
    return mean, std

train_transform = transforms.Compose([
        transforms.ToTensor(),
    ])

train_set = datasets.CIFAR10(root='./cifar_10', train=True, download=True, transform=train_transform)
train_loader = DataLoader(train_set, batch_size=100, shuffle=True)
mean, std = get_mean_std(train_loader)
print(mean)
print(std)

"""
Files already downloaded and verified
tensor([0.4914, 0.4822, 0.4465])
tensor([0.2470, 0.2435, 0.2616])
"""