compute_cifar10_mean_std.py

import numpy as np
from torchvision import datasets

train_dataset = datasets.CIFAR10("/tmp/data", train=True, download=True)
train_images = [np.array(train_dataset[i][0]) for i in range(len(train_dataset))]
train_images = np.stack(train_images)

means = (
    np.array(
        [
            np.mean(train_images[:, :, :, 0]),
            np.mean(train_images[:, :, :, 1]),
            np.mean(train_images[:, :, :, 2]),
        ]
    )
    / 255
)
stds = (
    np.array(
        [
            np.std(train_images[:, :, :, 0]),
            np.std(train_images[:, :, :, 1]),
            np.std(train_images[:, :, :, 2]),
        ]
    )
    / 255
)

print(means)
print(stds)