pytorch计算图像数据集的均值和方差

transforms.Normalize((0.485,0.456,0.406), (0.229,0.224,0.225))

# coding:utf-8
import os
import numpy as np
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms
from dataloader import Dataloader
from options import options
import pickle
"""
在训练前先运行该函数获得数据的均值和方差
"""

class Dataloader():def __init__(self, opt):# 训练&＃xff0c;验证&＃xff0c;测试数据集文件夹名self.opt &＃61; optself.dirs &＃61; [&＃39;train&＃39;, &＃39;test&＃39;, &＃39;testing&＃39;]self.means &＃61; [0, 0, 0]self.stdevs &＃61; [0, 0, 0]self.transform &＃61; transforms.Compose([transforms.Resize(opt.isize),transforms.CenterCrop(opt.isize),transforms.ToTensor(),#数据值从[0,255]范围转为[0,1]&＃xff0c;相当于除以255操作# transforms.Normalize((0.485,0.456,0.406), (0.229,0.224,0.225))])# 因为这里使用的是ImageFolder&＃xff0c;按文件夹给数据分类&＃xff0c;一个文件夹为一类&＃xff0c;label会自动标注好self.dataset &＃61; {x: ImageFolder(os.path.join(opt.dataroot, x), self.transform) for x in self.dirs}def get_mean_std(self, type, mean_std_path):"""
计算数据集的均值和方差:param type: 使用的是那个数据集的数据&＃xff0c;有&＃39;train&＃39;, &＃39;test&＃39;, &＃39;testing&＃39;:param mean_std_path: 计算出来的均值和方差存储的文件:return: """
num_imgs &＃61; len(self.dataset[type])for data in self.dataset[type]:img &＃61; data[0]for i in range(3):# 一个通道的均值和方差self.means[i] &＃43;&＃61; img[i, :, :].mean()self.stdevs[i] &＃43;&＃61; img[i, :, :].std()self.means &＃61; np.asarray(self.means) / num_imgsself.stdevs &＃61; np.asarray(self.stdevs) / num_imgsprint("{} : normMean &＃61; {}".format(type, self.means))print("{} : normstdevs &＃61; {}".format(type, self.stdevs))# 将得到的均值和方差写到文件中&＃xff0c;之后就能够从中读取with open(mean_std_path, &＃39;wb&＃39;) as f:pickle.dump(self.means, f)pickle.dump(self.stdevs, f)print(&＃39;pickle done&＃39;)if __name__ &＃61;&＃61; &＃39;__main__&＃39;:opt &＃61; options().parse()dataloader &＃61; Dataloader(opt)for x in dataloader.dirs:mean_std_path &＃61; &＃39;mean_std_value_&＃39; &＃43; x &＃43; &＃39;.pkl&＃39;dataloader.get_mean_std(x, mean_std_path)

# coding:utf-8
import os
import torch
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
import numpy as np
import pickle"""
用于加载训练train、验证test和测试数据testing
"""

class Dataloader():def __init__(self, opt):# 训练&＃xff0c;验证&＃xff0c;测试数据集文件夹名self.opt &＃61; optself.dirs &＃61; [&＃39;train&＃39;, &＃39;test&＃39;, &＃39;testing&＃39;]# 均值和方差存储的文件路径self.mean_std_path &＃61; {x: &＃39;mean_std_value_&＃39; &＃43; x &＃43; &＃39;.pkl&＃39; for x in self.dirs}# 初始化为0self.means &＃61; {x: [0, 0, 0] for x in self.dirs}self.stdevs &＃61; {x: [0, 0, 0] for x in self.dirs}print(type(self.means[&＃39;train&＃39;]))print(self.means)print(self.stdevs)for x in self.dirs:#如果存在则说明之前有获取过均值和方差if os.path.exists(self.mean_std_path[x]):with open(self.mean_std_path[x], &＃39;rb&＃39;) as f:self.means[x] &＃61; pickle.load(f)self.stdevs[x] &＃61; pickle.load(f)print(&＃39;pickle load done&＃39;)print(self.means)print(self.stdevs)# 将相应的均值和方差设置到transforms.Normalize函数中self.transform &＃61; {x: transforms.Compose([transforms.Resize(opt.isize),transforms.CenterCrop(opt.isize),transforms.ToTensor(),transforms.Normalize(self.means[x], self.stdevs[x]),]) for x in self.dirs}
...