Pytorch计算数据集的均值和方差

目录

文件夹的格式

代码

---

文件夹的格式

1. 没有分训练集测试集，就是一个文件夹下面分类别放。
2. '/home/ellery/workspace/OCLE/Plant-Pathology-master/data/images/'这个路径下有个casing_cap文件夹，这个文件夹下面有10个文件夹分别是10类，朋友们如果有三类，就可以对应大文件夹casing_cap下放置n个小的文件夹，即共有n个类别。
3. 我的网络层输入大小为512*512，所以transforms里先Resize设置成了这个尺寸，然后转成tensor，最后来计算所有数据的均值标准差。

代码

import os
import numpy as np
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms
import pickle

"""
在网络训练前先运行该函数获得数据的均值和标准差
"""

class Dataloader():
    def __init__(self, dataroot):
        self.dataroot = dataroot
        self.dirs = ['casing_cap']
        
        self.means = [0,0,0]
        self.std = [0,0,0]
        
        self.transform = transforms.Compose([transforms.Resize((512, 512)),
                                            transforms.ToTensor() # 数据值从[0,255]范围转为[0,1]，相当于除以255操作
                                            ])
        
        # 因为这里使用的是ImageFolder，按文件夹给数据分类，一个文件夹为一类，label会自动标注好
        self.dataset = {x: ImageFolder(os.path.join(dataroot, x), self.transform) for x in self.dirs}
        
    def get_mean_std(self):
        """
        计算数据集的均值和标准差
        """
        num_imgs = len(self.dataset['casing_cap'])
        for data in self.dataset['casing_cap']:
            img = data[0]
            for i in range(3):
                # 计算每一个通道的均值和标准差
                self.means[i] += img[i, :, :].mean()
                self.std[i] += img[i, :, :].std()
            
        self.means = np.asarray(self.means) / num_imgs
        self.std = np.asarray(self.std) /num_imgs
    
        print("{}: normMean = {}".format(type, self.means))
        print("{}: normstd = {}".format(type, self.std))
    
        # # 将得到的均值和标准差写到文件中，之后就能够从中读取
        # with open(mean_std_path, 'wb') as f:
        #     pickle.dump(self.means, f)
        #     pickle.dump(self.stdevs, f)
        #     print('pickle done')
        
if __name__ == '__main__':
    dataroot = '/home/ellery/workspace/OCLE/Plant-Pathology-master/data/images/'
    dataloader = Dataloader(dataroot)
    dataloader.get_mean_std()