6. 얼굴 이미지(HDF5 데이터 형식, GPU 가속)

 이번에는 GAN을 훈련해 사람의 얼굴을 생성해보겠다.

 

• CelebA 데이터셋

GAN 훈련에서 가장 큰 난관은 훈련하는 데에 피룡한 충분한 이미지를 확보하는 일이다. 사실 사람 이미지 10~100개만으로 GAN을 훈련할 수는 없다. 다행히도 굉장히 인기 있는 데이터셋인 CelebA 데이터셋을 이용할 수 있다.

 

• 계층적 데이터 형식

데이터를 특정 형식으로 바꾸면 좀 더 파일 접근일 용이하게 하고 반복에도 강하게 할 수 있다. 이 형식은 HDF5이다.

HDF5에서 5는 버전을 가리키며 HDF는 계층적 데이터 형식(Hierarchical Data Format)을 뜻한다. 이 포맷은 용량이 매우 큰 데이터에 효과적으로 접근하기 위해 만들어진 성숙한 데이터 형식으로, 과학이나 공학 응용 분야에서 널리 쓰인다.

 

HDF5 패키지가 계층적이라고 불리는 이유는, 하나 이상의 그룹을 가질 수 있기 때문이다. 또한 그룹 안에 여러 개의 데이터셋이 포함될 수도 있으며 그룹 안에 그룹이 존재하는 것도 가능하다. 이 방법은 우리에게 익숙한 폴더 구조와 비슷하다는 걸 알 수 있다.

• 데이터 가져오기

다음 코드는 CelebA 데이터셋을 다운로드 하고 20,000개 이미지를 추출하여 HDF5 파일로 묶는 코드이다.

 

import h5py

import numpy
import matplotlib.pyplot as plt
import zipfile
import imageio
import os
import torch

%%time

# location of the HDF5 package, yours may be under /gan/ not /myo_gan/
hdf5_file = '/content/drive/MyDrive/GAN-pytorch/celeba/celeba_aligned_small.h5py'

# how many of the 202,599 images to extract and package into HDF5
total_images = 20000

with h5py.File(hdf5_file, 'w') as hf:

    count = 0

    with zipfile.ZipFile('/content/drive/MyDrive/GAN-pytorch/celeba/img_align_celeba.zip', 'r') as zf:
      for i in zf.namelist():
        if (i[-4:] == '.jpg'):
          # extract image
          ofile = zf.extract(i)
          img = imageio.imread(ofile)
          os.remove(ofile)

          # add image data to HDF5 file with new name
          hf.create_dataset('img_align_celeba/'+str(count)+'.jpg', data=img, compression="gzip", compression_opts=9)
          
          count = count + 1
          if (count%1000 == 0):
            print("images done .. ", count)
            pass
            
          # stop when total_images reached
          if (count == total_images):
            break
          pass

        pass
      pass

 

 

• 데이터 살펴보기

h5py 라이브러리는 HDF5 패키지를 굉장히 파이써닉 하게 이용할 수 있게 해준다. 그룹이나 데이터셋을 파이썬 딕셔너리 객체처럼 사용하면 된다.

# open HDF5 file and list any groups

with h5py.File('/content/drive/MyDrive/GAN-pytorch/celeba/celeba_aligned_small.h5py', 'r') as file_object:
  
  for group in file_object:
    print(group)
    pass

 h5py 라이브러리로 HDF5 파일을 읽기 모드로 연다. 이는 파이썬에서 open()을 이용해서 파일을 여는 방식과 같다. 그다음 파일 객체를 순회하며 최상단에 있는 그룹의 이름을 출력한다.

 

 출력 결과 img_align_celeba라는 그룹 하나밖에 보이지 않는다. 그러면 파이썬에서 딕셔너리에 접근할 때와 같이 dataset=file_object['img_align_celeba']로 내부 데이터에 접근할 수 있다. 이를 통해 모든 이미지를 담은 데이터셋을 이용할 수 있다.

 

 물론 이 딕셔너리 형태의 문법을 이용해서 dataset['7.jpg']와 같이 개별적인 이미지에도 접근이 가능하다. 이미지 데이터 자체는 HDF5  형식으로 저장되어 있지만 넘파이 행렬로 쉽게 바꿀 수 있다.

# open HDF5 file, extract example image as numpy array and plot it

with h5py.File('/content/drive/MyDrive/GAN-pytorch/celeba/celeba_aligned_small.h5py', 'r') as file_object:
  dataset = file_object['img_align_celeba']
  image = numpy.array(dataset['6.jpg'])
  plt.imshow(image, interpolation='none')
  pass

 

 

• 데이터셋 클래스

class CelebADataset(Dataset):
  def __init__(self,file):
    self.file_object=h5py.File(file,'r')
    self.dataset=self.file_object['img_align_celeba']
    pass

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

  def __getitem__(self,index):
    if(index >= len(self.dataset)):
      raise IndexError()

    img=numpy.array(self.dataset[str(index) + '.jpg'])
    return torch.cuda.FloatTensor(img) / 255.0

  def plot_image(self,index):
    plt.imshow(numpy.array(self.dataset[str(index) + '.jpg']),interpolation='nearest')
    pass

  pass

 

 

• 판별기

MNIST 이미지와 CelebA 이미지가 유일하게 다른 점은 크기와 색깔 정보이다. CelebA 이미지는 218 x 178 픽셀 크기이고, 이는 218 x 178개 입력 노드가 있어야 한다는 뜻이다.

 

# discriminator class

class Discriminator(nn.Module):
    
    def __init__(self):
        # initialise parent pytorch class
        super().__init__()
        
        # define neural network layers
        self.model = nn.Sequential(
            View(218*178*3),
            
            nn.Linear(3*218*178, 100),
            nn.LeakyReLU(),
            
            nn.LayerNorm(100),
            
            nn.Linear(100, 1),
            nn.Sigmoid()
        )
        
        # create loss function
        self.loss_function = nn.BCELoss()

        # create optimiser, simple stochastic gradient descent
        self.optimiser = torch.optim.Adam(self.parameters(), lr=0.0001)

        # counter and accumulator for progress
        self.counter = 0;
        self.progress = []

        pass
    
    
    def forward(self, inputs):
        # simply run model
        return self.model(inputs)
    
    
    def train(self, inputs, targets):
        # calculate the output of the network
        outputs = self.forward(inputs)
        
        # calculate loss
        loss = self.loss_function(outputs, targets)

        # increase counter and accumulate error every 10
        self.counter += 1;
        if (self.counter % 10 == 0):
            self.progress.append(loss.item())
            pass
        if (self.counter % 1000 == 0):
            print("counter = ", self.counter)
            pass

        # zero gradients, perform a backward pass, update weights
        self.optimiser.zero_grad()
        loss.backward()
        self.optimiser.step()

        pass
    
    
    def plot_progress(self):
        df = pandas.DataFrame(self.progress, columns=['loss'])
        df.plot(ylim=(0), figsize=(16,8), alpha=0.1, marker='.', grid=True, yticks=(0, 0.25, 0.5, 1.0, 5.0))
        pass
    
    pass

 

 View는 3차원 이미지 텐서를 1차원 형태의 텐서로 바꿔주는 역할을 한다.

class View(nn.Module):
  def __init__(self,shape):
    super().__init__()
    self.shape=shape,
    

  def forward(self,x):
    return x.view(*self.shape)

 

• 판별기 테스트하기

  %%time
  # test discriminator can separate real data from random noise
  
  D = Discriminator()
  # move model to cuda device
  D.to(device)
  
  for image_data_tensor in celeba_dataset:
      # real data
      D.train(image_data_tensor, torch.cuda.FloatTensor([1.0]))
      # fake data
      D.train(generate_random_image((218,178,3)), torch.cuda.FloatTensor([0.0]))
      pass​

 

• GPU 가속

# check if CUDA is available
# if yes, set default tensor type to cuda

if torch.cuda.is_available():
  torch.set_default_tensor_type(torch.cuda.FloatTensor)
  print("using cuda:", torch.cuda.get_device_name(0))
  pass

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

device

 

 generate_random_image() 와 generate_random_seed() 함수는 굳이 수정 할 필요는 없다. 기본 float 탕비이 이제 CUDA 텐서로 처리될 것이기 때문이다.

 

 

• 생성기

3차원의 텐서를 (218,178,3) 크기로 결과를 내어주도록 바꾸어야 한다.

class Generator(nn.Module):
    
    def __init__(self):
        # initialise parent pytorch class
        super().__init__()
        
        # define neural network layers
        self.model = nn.Sequential(
            nn.Linear(100, 3*10*10),
            nn.LeakyReLU(),
            
            nn.LayerNorm(3*10*10),
            
            nn.Linear(3*10*10, 3*218*178),
            
            nn.Sigmoid(),
            View((218,178,3))
        )
        
        # create optimiser, simple stochastic gradient descent
        self.optimiser = torch.optim.Adam(self.parameters(), lr=0.0001)

        # counter and accumulator for progress
        self.counter = 0;
        self.progress = []
        
        pass
    
    
    def forward(self, inputs):        
        # simply run model
        return self.model(inputs)
    
    
    def train(self, D, inputs, targets):
        # calculate the output of the network
        g_output = self.forward(inputs)
        
        # pass onto Discriminator
        d_output = D.forward(g_output)
        
        # calculate error
        loss = D.loss_function(d_output, targets)

        # increase counter and accumulate error every 10
        self.counter += 1;
        if (self.counter % 10 == 0):
            self.progress.append(loss.item())
            pass

        # zero gradients, perform a backward pass, update weights
        self.optimiser.zero_grad()
        loss.backward()
        self.optimiser.step()

        pass
    
    
    def plot_progress(self):
        df = pandas.DataFrame(self.progress, columns=['loss'])
        df.plot(ylim=(0), figsize=(16,8), alpha=0.1, marker='.', grid=True, yticks=(0, 0.25, 0.5, 1.0, 5.0))
        pass
    
    pass

 

• 생성기 결과 확인

G=Generator()

G.to(device)

output=G.forward(generate_random_seed(100))
img=output.detach().cpu().numpy()
plt.imshow(img,interpolation='none',cmap='Blues')

 이제 출력값을 실제 이미지로 그려내기 전에, detach() 를 써서 파이토치의 계산 그래프로부터 값을 떼어내 GPU 에서 CPU로 옮기고 넘파이 행렬로 변환한다.

 

 출력을 확인해보면 올바른 크기이고 임의의 색으로 채워진 이미지를 확인할 수 있다. 이미지가 거의 하나의 색깔로 채워져 있거나 어떤 패턴이 보인다면 코드에 뭔가 문제가 있다고 의심해봐야 한다. 훈련되지 않은 생성기는 단지 아무 패턴이 없는 임의의 데이터를 만들어내야 한다.

 

• GAN 훈련

 판별기와 생성기를 GPU로 옮겨주고 목푯값을 torch.cuda.FloatTensor 타입으로 변경한다.

%%time

# 판별기 및 생성기 생성

D=Discriminator()
D.to(device)
G=Generator()
G.to(device)

epochs=1

for epoch in range(epochs):
  print('epoch=',epoch + 1)

  # 판별기와 생성기 훈련
  for image_data_tensor in celeba_dataset:
    # 참에 대해 판별기 훈련
    D.train(image_data_tensor,torch.cuda.FloatTensor([1.0]))

    # 거짓에 대해 판별기 훈련
    # G의 기울기가 계산되지 않도록 detach() 함수를 이용
    D.train(G.forward(generate_random_seed(100)).detach(),torch.cuda.FloatTensor([0.0]))

    # 생성기 훈련
    G.train(D,generate_random_seed(100),torch.cuda.FloatTensor([1.0]))
    pass

  pass

# plot several outputs from the trained generator

# plot a 3 column, 2 row array of generated images
f, axarr = plt.subplots(2,3, figsize=(16,8))
for i in range(2):
    for j in range(3):
        output = G.forward(generate_random_seed(100))
        img = output.detach().cpu().numpy()
        axarr[i,j].imshow(img, interpolation='none', cmap='Blues')
        pass
    pass

 

 사실 이러한 이미지의 다양성에 대해 한 가지 의문이 생기긴 한다. 생성기가 정말로 무언가를 배운 것일까? 생성기는 직접 이미지를 보고 배운 것이 아니므로 얼굴의 전체적인 형태나 눈, 코 등의 세세한 특징을 기억하지 못한다. 사실 생성기가 배운 것은 이미지를 생성할 때 훈련 데이터의 우도(가능도)(likelihood)이다.