캠프60일차 - 이미지 딥러닝 응용(Style Transfer Learning, GAN)

GAN을 이용하여 얼굴을 학습하고, 생성하는 모델을 만들어본다.

생성적 적대 신경망(GAN, Generative Adversarial Network)이란?

두 개의 신경망(생성자와 판별자)이 서로 경쟁하며 학습하는 구조를 가진다. 새로운 데이터를 생성하거나 기존 데이터를 변형하는 데 사용되며, 이미지 생성, 데이터 증강, 이상 탐지 등에 활용된다.

생성자 (Generator)는 랜덤한 노이즈를 바탕으로 가짜 데이터를 만들며, 판별자를 속일 수 있을 정도의 진짜 같은 데이터를 생성하는데 목적이 있다.

판별자 (Discriminator)는 입력받은 데이터가 실제(real) 데이터인지 생성된(fake) 데이터인지 구분한다. 목표는 생성자가 만든 가짜 데이터를 정확히 구별하는 것이다.

두 네트워크는 서로 경쟁적으로 학습하며, 이를 통해 생성자는 점점 더 진짜 같은 데이터를 생성하고, 판별자는 점점 더 정교하게 가짜 데이터를 구분하기 때문에 적대적 학습(adversarial training)이라고 부른다.

!pip install -q kagglehub
import kagglehub

# Download latest version
path = kagglehub.dataset_download("jessicali9530/celeba-dataset")

print("Path to dataset files:", path)

kaggle에 있는 얼굴 이미지 데이터를 가져온다.

이미지를 보면 다음과 같이 사람 얼굴 부분에 해당하는 이미지 셋임을 알 수 있다.

데이터를 모델에 넣기 위해 전처리를 진행한다.

import torch
import torchvision.transforms as tf
from torchvision.datasets import ImageFolder
from torch.utils.data.dataloader import DataLoader

transforms = tf.Compose([
    tf.Resize(64),
    tf.CenterCrop(64),
    tf.ToTensor(),
    tf.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = ImageFolder(
    root = dir + "/img_align_celeba/",
    transform=transforms
)
loader = DataLoader(dataset, batch_size = 128, shuffle=True)

torchvision의 Compose를 이용해 이미지의 크기를 맞춰준다. 이 때 Resize와 CenterCrop을 각각 해주게 되면 이미지의 비율을 유지하면서 원하는 크기로 이미지를 잘라 가져올 수 있다. Normalize는 앞은 타겟으로 하는 평균, 뒤는 타겟으로 하는 표준편차에 맞추어 그림 데이터를 Normalize하는 역할을 한다.

import torch.nn as nn
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        # 생성자를 구성하는 총 정의
        self.gen = nn.Sequential(
            nn.ConvTranspose2d(100,512,kernel_size=4, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(),
            
			...
            
            nn.ConvTranspose2d(64,3,kernel_size=4,stride=2,padding=1,bias=False),
            nn.Tanh()
        )
    def forward(self, x):
        return self.gen(x)

다음과 같이 생성자(Generator) class를 만들어 준다. ConvTranspose2d - BatchNorm2d - ReLU 블럭으로 구성되며 마지막에만 Tanh()를 통과한다.

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        # 감별자 구성 층 정의
        self.disc = nn.Sequential(
            nn.Conv2d(3,64,kernel_size=4,stride=2,padding=1,bias=False),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2),
            
			...
            
            nn.Conv2d(512,1,kernel_size=4),
            nn.Sigmoid()
            
    )
    def forward(self, x):
        return self.disc(x)

생성자 데이터의 참거짓을 판단하는 감별자(Discriminator)를 선언한다. 참/거짓을 분류하므로 Sigmoid를 통과한다.

# GAN의 학습을 원활하게 이루어지도록 가중치 초기화 설정
def weights_init(m):
    # 층의 종류 추출
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        # 합성곱 초기화
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        # 배치 정규화층 초기화
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

학습을 원활하게 할 수 있도록 가중치를 초기화 해준다.

import tqdm
from torch.optim.adam import Adam

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# 생성자 정의
G = Generator().to(device)
G.apply(weights_init)

# 감별자 정의
D = Discriminator().to(device)
D.apply(weights_init)

G_optim = Adam(G.parameters(),lr=0.0001, betas=(0.5, 0.999))
D_optim = Adam(D.parameters(),lr=0.0001, betas=(0.5, 0.999))

best_loss1 = float('inf')
best_loss2 = float('inf')

Adam 옵티마이져를 생성자와 감별자에 대해 각각 정의해준다. loss 값을 저장하기 위해 무한대로 초기화한다.

for epochs in range(5):
    iterator = tqdm.tqdm(enumerate(loader,0),total=len(loader))
    
    for i, data in iterator:
        D_optim.zero_grad()
        # 실제 이미지 1, 생성 이미지 0으로 정답 설정
        label = torch.ones_like(data[1], dtype=torch.float32).to(device)
        label_fake = torch.zeros_like(data[1], dtype=torch.float32).to(device)
        # 실제 이미지를 감별자에 입력
        real = D(data[0].to(device))
        # 실제 이미지에 대한 감별자의 오차 계산
        Dloss_real = nn.BCELoss()(torch.squeeze(real), label)
        Dloss_real.backward()
        
        # 감별자 학습
        # 가짜 이미지 생성
        noise = torch.randn(label.shape[0], 100, 1, 1, device=device)
        fake = G(noise)
        # 가짜 이미지를 감별자에 입력
        output = D(fake.detach())
        
        # 가짜 이미지에 대한 감별자의 오차 계산
        Dloss_fake = nn.BCELoss()(torch.squeeze(output), label_fake)
        Dloss_fake.backward()
        
        # 감별자의 전체 오차를 학습
        Dloss = 0.5 * (Dloss_real + Dloss_fake)
        D_optim.step()
        
        # 생성자의 학습
        G_optim.zero_grad()
        output = D(fake)
        Gloss = nn.BCELoss()(torch.squeeze(output),label)
        Gloss.backward()
        
        G_optim.step()
        
        iterator.set_description(f'epoch:{epochs} iteration:{i+1} D_loss:{Dloss} Gloss:{Gloss}')
        
        # if Gloss.item() < best_loss1:
        #     best_loss1 = Gloss.item()
        #     torch.save(G.state_dict(), 'Generator.pth')
        #     torch.save(D.state_dict(), 'Discriminator.pth')
        #     print(f"Best model saved at epoch {epochs}, batch {i} with G loss: {Gloss.item()} and D loss: {Dloss.item()}")
    
    if Gloss.item() < best_loss2:
          best_loss2 = Gloss.item()
          torch.save(G.state_dict(), 'Generator2.pth')
          torch.save(D.state_dict(), 'Discriminator2.pth')
          print(f"Best model saved at epoch {epochs}, batch {i} with G loss: {Gloss.item()} and D loss: {Dloss.item()}")

torch.save(G.state_dict(), 'Generator3.pth')
torch.save(D.state_dict(), 'Discriminator3.pth')

다음과 같이 감별자는 실제 이미지를 1, 생성자가 생성한 이미지를 0으로 판단하도록 정답을 설정한다. 가짜 이미지에 대해 감별자의 오차를 계산한다. 생성자는 가짜 이미지를 생성하고 감별자의 오차를 바탕으로 역전파를 통해 학습을 한다. 이를 통해 감별자는 가짜 이미지를 판단하는 학습을, 생성자는 감별자를 속일 수 있는 수준의 이미지를 만든다.

with torch.no_grad():
    G.load_state_dict(torch.load('Generator2.pth',map_location=device))
    D.load_state_dict(torch.load('Discriminator2.pth',map_location=device))
    # 랜덤한 하나의 점 지정
    feature_vector = torch.randn(1, 100, 1, 1).to(device)
    
    pred = G(feature_vector).squeeze()
    pred = pred.permute(1,2,0).cpu().numpy()
    
    plt.imshow(pred)
    plt.title('predicted image')
    plt.show()

GAN을 출력하기 위해 값의 입력이 필요하기 때문에 임의의 값을 넣어준다.

5번 학습한 데이터 중 잘 나온 데이터 선택 (epoch 당 G loss로 선택한 모델, 마지막 학습 모델)

50번 학습한 데이터 중 잘 나온 데이터 선택 (50번 후 저장한 모델, 50번 중 Gloss 기준 저장한 모델)

해당 구조는 5번 정도 학습이 이루어졌을 때 이미 loss의 변화가 크게 일어나지 않았다. 모델을 불러와 이미지를 생성하였으며 대부분 예상하는 만큼 얼굴을 잘 표현하지는 못하며 그 중 보기에 잘 나오는 이미지를 선택하였다.

첫번째 그림은 각 epoch 당 loss를 비교하여 loss가 감소하면 저장하는 방식을 채택하였다. 이 방식이 모델의 학습 가중치의 편향을 일으키지 않으면서 최적의 값을 가질 수 있을 것이라 생각하였다. 2번째 학습을 마치고 저장한 모델이다. 두번째 그림은 5번을 전부 학습한 다음 도출한 그림이다. loss가 다시 높게 판별 되는 경우를 무시하고 저장하기 때문에 최적의 값을 가지기 어려울 수 있다고 생각된다. 세번째 그림은 epoch 안에서 판별 조건을 설정하여 저장한 경우이다. 이 경우 우연히 D loss가 클 때 (판별을 잘 못할 때) G loss가 크게 감소하여 저장이 이루어지는 것을 확인하였다. epoch 안에서 모델을 선택하는 것은 특수한 경우를 선택할 가능성이 높으므로 올바르게 모델을 저장했다고 보기 어렵다고 판단된다. 마지막 그림은 50번 학습을 마친 뒤 생성한 모습이다. 2번째 그림과 마찬가지로 최적의 경우를 고르기 어려울 것 같다고 판단된다.

loss값을 관찰하여 얻은 가설 및 판단은 각 epoch 당 loss를 비교하는 것이 특수한 경우에 빠지는 것을 방지하기 조금 더 유리함과 동시에 그 중 최적의 값을 고를 수 있을 것이라는 점이다. 물론 이렇게 선택한 경우에도 loss가 튀는 경우가 발생할 수 있기 때문에 모델 학습 시 저장하는 위치에 대한 고민이 더 필요해 보인다.