[PyTorch] 이미지 픽셀의 평균, 표준편차를 계산하여 정규화하기

 dataset에 있는 이미지의 평균과 표준편차를 계산하여 정규화(normalize) 해보겠습니다.

 

1. 데이터셋을 불러옵니다.

 저는 torchvision에서 제공하는 STL-10 dataset을 사용했습니다. train dataset을 불러와서 train_ds에 저장합니다.

 

# loading training data
from torchvision import datasets
import torchvision.transforms as transforms
import os

path2data = '/data'

# if not exists the path, make the path
if not os.path.exists(path2data):
    os.mkdir(path2data)

data_transformer = transforms.Compose([transforms.ToTensor()])
train_ds = datasets.STL10(path2data, split='train', download='True', transform=data_transformer)

print(train_ds.data.shape)

 

 

2. train_ds의 평균과 표준편차를 계산합니다.

# calculate the mean and standard deviation of train_ds
import numpy as np

meanRGB = [np.mean(x.numpy(), axis=(1,2)) for x,_ in train_ds]
stdRGB = [np.std(x.numpy(), axis=(1,2)) for x,_ in train_ds]

meanR = np.mean([m[0] for m in meanRGB])
meanG = np.mean([m[1] for m in meanRGB])
meanB = np.mean([m[2] for m in meanRGB])

stdR = np.mean([s[0] for s in stdRGB])
stdG = np.mean([s[1] for s in stdRGB])
stdB = np.mean([s[2] for s in stdRGB])

print(meanR, meanG, meanB)
print(stdR, stdG, stdB)

 

 

3. Normalize transformation을 적용합니다.

Normalize 이외에도 다른 transformation을 적용했습니다.

train_transformer = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([meanR, meanG, meanB], [stdR, stdG, stdB])
])

train_ds.transform = train_transformer

 

4. transformation이 적용된 이미지를 확인합니다.

# display the transformed sample images from train_ds
import torch
np.random.seed(0)
torch.manual_seed(0)

grid_size = 4
rnd_inds = np.random.randint(0, len(train_ds), grid_size)
print('image indices:', rnd_inds)

x_grid = [train_ds[i][0] for i in rnd_inds]
y_grid = [train_ds[i][1] for i in rnd_inds]

x_grid = utils.make_grid(x_grid, nrow=4, padding=2)
print(x_grid.shape)

plt.figure(figsize=(10,10))
show(x_grid, y_grid)