在深度学习中计算自定义数据集的均值和标准差

一、前言：对于常用的图像分类模型，如 VGG、ResNet 等，这些值通常是根据 ImageNet 数据集统计得出的。ImageNet 数据集包含了大量的自然图像，通过对这些图像进行统计分析，得到了如下均值和标准差：mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)。但是对于自己的数据集，想知道其均值和标准差，又该如何计算呢？

二、下面使用Pytorch框架来实现计算自己的数据集的均值和标准差

1.使用CustomDataset类读取数据集，遍历数据集中的所有图像（不需要标签）

class CustomDataset(Dataset):
    def __init__(self, root_images, transform=None):
        self.root_images = root_images
        self.transform = transform
        self.files = sorted(os.listdir(root_images))

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        img_path = os.path.join(self.root_images, self.files[idx])
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image

2.定义计算均值和标准差的函数

def compute_mean_std(root_images, batch_size=32, num_workers=4):
    transform = transforms.Compose([
        transforms.Resize((512, 512)),  # 根据需要调整尺寸
        transforms.ToTensor()
    ])

    dataset = CustomDataset(root_images, transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

3.初始化均值（mean）和标准差（std）

    mean = 0.
    std = 0.

4.使用for循环遍历数据集，计算均值标准差，并返回

    # 遍历数据集
    for images in dataloader:  # 修改这里，不要尝试解包两个值
        batch_samples = images.size(0)  # 获取当前批次的样本数量
        images = images.view(batch_samples, images.size(1), -1)
        mean += images.mean(2).sum(0)
        std += images.std(2).sum(0)

    # 计算均值和标准差
    mean /= len(dataset)
    std /= len(dataset)

    return mean.numpy(), std.numpy()

三、最终完整版代码如下：在运行前先安装所需要的包

import os
import torch
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
import numpy as np

# 定义一个数据集类
class CustomDataset(Dataset):
    def __init__(self, root_images, transform=None):
        self.root_images = root_images
        self.transform = transform
        self.files = sorted(os.listdir(root_images))

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        img_path = os.path.join(self.root_images, self.files[idx])
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image

def compute_mean_std(root_images, batch_size=32, num_workers=4):
    transform = transforms.Compose([
        transforms.Resize((512, 512)),  # 根据需要调整尺寸
        transforms.ToTensor()
    ])

    dataset = CustomDataset(root_images, transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

    # 初始化均值和方差
    mean = 0.
    std = 0.

    # 遍历数据集
    for images in dataloader:  # 修改这里，不要尝试解包两个值
        batch_samples = images.size(0)  # 获取当前批次的样本数量
        images = images.view(batch_samples, images.size(1), -1)
        mean += images.mean(2).sum(0)
        std += images.std(2).sum(0)

    # 计算均值和标准差
    mean /= len(dataset)
    std /= len(dataset)

    return mean.numpy(), std.numpy()

if __name__ == '__main__':
    root_images = r'原始图像的路径'
    mean, std = compute_mean_std(root_images)
    print("Mean:", mean)
    print("Std:", std)

以上内容如果对你有帮助的话，请点个赞吧，谢谢！