Pytorch图像预处理时,通常使用transforms.Normalize(mean, std)对图像按通道进行标准化,即减去均值,再除以方差。这样做可以加快模型的收敛速度。其中参数mean和std分别表示图像每个通道的均值和方差序列。
Imagenet数据集的均值和方差为:mean=(0.485, 0.456, 0.406),std=(0.229, 0.224, 0.225),因为这是在百万张图像上计算而得的,所以我们通常见到在训练过程中使用它们做标准化。而对于特定的数据集,选择这个值的结果可能并不理想。接下来给出计算特定数据集的均值和方差的方法。
def getStat(train_data):
'''
Compute mean and variance for training data
:param train_data: 自定义类Dataset(或ImageFolder即可)
:return: (mean, std)
'''
print('Compute mean and variance for training data.')
print(len(train_data))
train_loader = torch.utils.data.DataLoader(
train_data, batch_size=1, shuffle=False, num_workers=0,
pin_memory=True)
mean = torch.zeros(3)
std = torch.zeros(3)
for X, _ in train_loader:
for d in range(3):
mean[d] += X[:, d, :, :].mean()
std[d] += X[:, d, :, :].std()
mean.div_(len(train_data))
std.div_(len(train_data))
return list(mean.numpy()), list(std.numpy())
if __name__ == '__main__':
train_dataset = ImageFolder(root=r'D:\cifar10_images\test', transform=None)
print(getStat(train_dataset))
getState()方法接收一个Dataset类(ImageFolder),然后累加所有图像三个通道的均值和方差,最后除以图像总数并返回。
这里用cifar10做的测试,测试集返回的结果如下所示:
Compute mean and variance for training data.
10000
([0.4940607, 0.4850613, 0.45037037], [0.20085774, 0.19870903, 0.20153421])