고양이 이미지를 학습해서 새로운 고양이 이미지를 만들어주는 GAN을 제작해보자.
만들기 전에 GAN에 대한 설명과 예시는 아래 링크를 참조하자.
이 글에서 catGAN은 카테고리 GAN이 아니라 고양이 GAN이다.
1. 데이터 준비
Kaggle에서 데이터를 받아올 것이다.
이 링크를 통해 이미지를 다운받거나 아래 명령어를 터미널에 복사해서 붙여 넣고 실행하자.
curl -L -o ~/Downloads/archive.zip https://www.kaggle.com/api/v1/datasets/download/spandan2/cats-faces-64x64-for-generative-models
cURL을 사용하려 할 때, Kaggle API가 설정되어있지 않다면
- Kaggle에 로그인한 뒤 My Account 페이지에서
- 페이지 하단에 Create New API Token을 클릭해서
kaggle.json
파일을 다운로드하고 - 이 파일을 홈 디렉토리 아래
.kaggle
폴더에 복사해넣는다. - 보안을 위해 터미널에서
chmod 600 ~/.kaggle/kaggle.json
을 실행해서 파일 권한을 설정한다.
이제 받은 파일의 압축을 푼다.
나는 data/cats/
에 압축을 풀었다.
2. Generator 및 Discriminator 모델 설계
이제 모델을 설계할 차례이다.
편의상 주피터 랩에서 코드를 작성했다.
먼저 필요한 라이브러리와 함수를 임포트한다.
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
이제 generator와 discriminator을 설계하자.
# Generator 모델 정의
class Generator(nn.Module):
def __init__(self, nz, ngf, nc):
super(Generator, self).__init__()
self.main = nn.Sequential(
nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),
nn.BatchNorm2d(ngf * 8),
nn.ReLU(True),
nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf * 4),
nn.ReLU(True),
nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf * 2),
nn.ReLU(True),
nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf),
nn.ReLU(True),
nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
nn.Tanh()
)
def forward(self, input):
return self.main(input)
# Discriminator 모델 정의
class Discriminator(nn.Module):
def __init__(self, nc, ndf):
super(Discriminator, self).__init__()
self.main = nn.Sequential(
nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 2),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 4),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 8),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
nn.Sigmoid()
)
def forward(self, input):
return self.main(input)
참고로 64*64 이미지를 처리하는데 널리 사용되는 지극히 평범한 모델이다.
3. 모델 초기화, 손실 함수 및 최적화 기법 설정
다운받은 고양이 이미지를 가져오는 코드와 모델을 초기화하고, 각 모델의 가중치를 초기화하는 코드를 만들자.
# 데이터셋 로더 설정
def load_data(batch_size):
dataset = ImageFolder(
root="data/cats",
transform=transforms.Compose([
transforms.Resize(64),
transforms.CenterCrop(64),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
return dataloader
이 코드는 고양이 이미지 데이터셋을 로드하여 DataLoader
형태로 반환한다.
이미지를 64x64 크기로 조정하고 (0.5, 0.5, 0.5)
평균 및 표준편하로 정규화 하여 모델 학습에 맞게 변환한다.
데이터를 저장한 곳으로 경로를 잘 설정하자.
# 모델 초기화
def initialize_models(nz, ngf, ndf, nc, device):
netG = Generator(nz, ngf, nc).to(device)
netD = Discriminator(nc, ndf).to(device)
netG.apply(weights_init)
netD.apply(weights_init)
return netG, netD
# 가중치 초기화 함수
def weights_init(m):
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
nn.init.normal_(m.weight.data, 0.0, 0.02)
elif isinstance(m, nn.BatchNorm2d):
nn.init.normal_(m.weight.data, 1.0, 0.02)
nn.init.constant_(m.bias.data, 0)
이 코드에서 initialize_models
함수는 Generator와 Discriminator 모델을 초기화하고, 각 모델의 가중치를 weights_init
함수로 설정한 수 반환하는 코드이다.
weights_init
함수는 Convolution 및 BatchNorm 레이어의 가중치를 초기화하는 함수로 각 모델의 초기 가중치를 설정해서 학습의 안전성을 높인다.
Convolution 레이어는 평균이 0 표준편차가 0.02인 분포로 초기화되고 BatchNorm 레이어는 평균이 1 표준 편자가 0.02인 정규분포로 초기화된다.
4. 학습
이제 모델을 트레이닝하는 함수를 만든다.
fixed_noise
을 설정해서 고정된 latent vector을 사용해 학습 과정 동안 생성된 이미지가 어떻게 변하는지 관찰해보자.
# 학습 함수
def train_dcgan(num_epochs, dataloader, netG, netD, device, fixed_noise, save_interval=5):
criterion = nn.BCELoss()
optimizerD = optim.Adam(netD.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=0.0002, betas=(0.5, 0.999))
for epoch in range(num_epochs):
for i, (real_images, _) in enumerate(dataloader):
# Discriminator 학습
netD.zero_grad()
real_images = real_images.to(device)
b_size = real_images.size(0)
label = torch.full((b_size,), 1., device=device)
output = netD(real_images).view(-1)
errD_real = criterion(output, label)
errD_real.backward()
D_x = output.mean().item()
noise = torch.randn(b_size, nz, 1, 1, device=device)
fake_images = netG(noise)
label.fill_(0.)
output = netD(fake_images.detach()).view(-1)
errD_fake = criterion(output, label)
errD_fake.backward()
D_G_z1 = output.mean().item()
optimizerD.step()
# Generator 학습
netG.zero_grad()
label.fill_(1.)
output = netD(fake_images).view(-1)
errG = criterion(output, label)
errG.backward()
D_G_z2 = output.mean().item()
optimizerG.step()
# 학습 진행 출력
if i % 50 == 0:
print(f"[{epoch}/{num_epochs}][{i}/{len(dataloader)}] "
f"Loss_D: {errD_real.item() + errD_fake.item():.4f} "
f"Loss_G: {errG.item():.4f} "
f"D(x): {D_x:.4f} D(G(z)): {D_G_z1:.4f}/{D_G_z2:.4f}")
# 고정된 노이즈로 결과 이미지 생성 및 모델 저장
if not os.path.exists("results"):
os.makedirs("results")
if not os.path.exists("models"):
os.makedirs("models")
if epoch % save_interval == 0 or epoch == num_epochs - 1:
with torch.no_grad():
fake_images_fixed = netG(fixed_noise).detach().cpu()
# 비정규화: [-1, 1] 범위를 [0, 1]로 변환
fake_images_fixed = (fake_images_fixed + 1) / 2
vutils.save_image(fake_images_fixed, f"results/fake_samples_epoch_{epoch}.png")
torch.save(netG.state_dict(), f"models/netG_epoch_{epoch}.pth")
torch.save(netD.state_dict(), f"models/netD_epoch_{epoch}.pth")
위 함수는 전체 학습 과정을 담당하는 코드로 설명은 맨 위 링크의 GAN의 학습 과정을 참고하자.
이제 하이퍼파라미터들을 설정하고 학습을 진행해보자.
latent vector의 사이즈는 100, 배치 사이즈는 128로 두고 50 에포크 동안 학습시켜보자.
애플실리콘에서도 잘 돌아갈 수 있도록 mps도 사용할 수 있게 했다.
# 주요 하이퍼파라미터 설정
device = torch.device('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
nz = 100 # latent vector size
ngf = 64 # size of feature maps in generator
ndf = 64 # size of feature maps in discriminator
nc = 3 # number of color channels
batch_size = 128
num_epochs = 50
save_interval = 5
# 학습 프로세스 실행
dataloader = load_data(batch_size)
netG, netD = initialize_models(nz, ngf, ndf, nc, device)
fixed_noise = torch.randn(64, nz, 1, 1, device=device)
train_dcgan(num_epochs, dataloader, netG, netD, device, fixed_noise, save_interval)
이 코드를 실행하면 아래와 같은 출력이 나온다.
출력의 각 값의 의미를 짚고 가자.
내 코드에서 마지막 줄의 출력은 아래와 같다.
[49/50][100/124] Loss_D: 0.4689 Loss_G: 2.6872 D(x): 0.9234 D(G(z)): 0.2776/0.1053
이 출력은 각 모델의 학습 상태를 나타내는 로그이다.
Loss_D (0.4689)
- Loss_D는 Discriminator의 손실값이다.
- Discriminator는 "진짜 이미지"와 "가짜 이미지"를 판별하는 역할을 하므로, 손실값이 낮을수록 Discriminator가 더 잘 판별하고 있다는 의미이다.
Loss_D = 0.4689
는 Discriminator가 가짜와 진짜 이미지를 꽤 잘 구별하고 있다는 것을 의미한다.
Loss_G (2.6872)
- Loss_G는 Generator의 손실값이다.
- Generator는 Discriminator가 가짜 이미지를 진짜로 판별하도록 학습하는데, 이 손실값이 낮을수록 Generator가 더 진짜 같은 이미지를 생성하고 있다는 의미이다.
Loss_G = 2.6872
는 Generator가 아직 Discriminator를 속이기 위해 더 많이 학습해야 한다는 것을 나타낸다.
D(x) (0.9234)
- D(x)는 Discriminator가 "진짜 이미지(x)"에 대해 예측한 확률이다.
- Discriminator가 진짜 이미지를 얼마나 잘 분별하는지 나타낸다. 값이 1에 가까울수록 Discriminator는 진짜 이미지를 잘 판별한다.
D(x) = 0.9234
는 Discriminator가 진짜 이미지를 92.34% 확률로 진짜라고 판단한 것에 해당한다.
D(G(z)) (0.2776)
- D(G(z))는 Discriminator가 Generator가 생성한 가짜 이미지(G(z))에 대해 예측한 확률이다.
- Discriminator가 가짜 이미지를 얼마나 잘 판별하는지 나타낸다. 값이 0에 가까울수록 Discriminator가 가짜 이미지를 잘 판별한 것이다.
D(G(z)) = 0.2776
은 Discriminator가 Generator가 생성한 가짜 이미지를 27.76% 확률로 진짜라고 판단했다는 의미이다. 이는 Generator가 Discriminator를 속이기 위해 더 많이 학습해야 한다는 것을 뜻한다.
D(G(z)) (0.1053)
- D(G(z))의 두 번째 값은 현재 Generator가 만든 가짜 이미지에 대한 판별 결과이다. 학습이 진행됨에 따라 이 값이 점차 낮아져야, Generator가 더 진짜 같은 이미지를 생성하고 있다는 뜻이다.
D(G(z)) = 0.1053
는 Generator가 만든 가짜 이미지에 대해 Discriminator가 10.53% 확률로 진짜로 판별했다는 뜻이다. 이 값이 낮을수록 Generator가 Discriminator를 잘 속였다는 의미다.
요약:
- Loss_D는 Discriminator의 학습 상태, Loss_G는 Generator의 학습 상태를 나타낸다.
- D(x)는 Discriminator가 진짜 이미지를 판별하는 능력을 나타내고, D(G(z))는 Generator가 만든 가짜 이미지를 판별하는 능력을 나타낸다.
5. 결과 확인
result
폴더에서 가장 처음 파일과 마지막 파일을 보자.
꽤 그럴듯한 고양이들이 생겼다.
학습 과정에서 저장된 이미지들을 모아 gif 파일로 만들어봤다.
그리고 훈련 시간을 대폭 키우고 좀 예쁘게 생긴 고양이들이 나오는 노이즈만 모아서 훈련에 따른 결과의 변화를 살펴보았다.