1207
[16주차 - Day4] Monthly Project
실전 프로젝트 - GAN을 이용한 얼굴 마스크 이미지 생성
# 팀원분들 불참 & 내 기말고사가 겹쳐서 실습 베이직 코드만 정리
데이터셋 다운로드
실습을 위해 Face Mask Classification 데이터셋
데이터셋은 흑백(grayscale) 이미지로 구성됩니다.
GPU 가속기를 이용하기 위해 [런타임] - [런타임 유형 변경] - [GPU] 를 선택합니다.
학습 데이터셋
마스크 착용(with mask): 9,000장
마스크 미착용(wihtout mask): 9,000장
테스트 데이터셋
마스크 착용(with mask): 1,000장
마스크 미착용(wihtout mask): 1,000장
!git clone https://github.com/ndb796/Face-Mask-Classification-20000-Dataset
필요한 라이브러리 불러오기
실습을 위한 PyTorch 라이브러리를 불러옵니다.
import os
import time
import matplotlib.pyplot as plt
from PIL import Image
import torch
import torch.nn as nn
from torchvision import datasets
import torchvision.transforms as transforms
from torchvision.utils import save_image
학습 이미지 시각화
학습할 이미지 중에서 원하는 이미지를 골라 시각화합니다.
# 마스크를 쓴 이미지 출력
image = Image.open('./Face-Mask-Classification-20000-Dataset/train/with_mask/with-mask-default-mask-seed0030.png')
print("이미지 크기:", image.size)
plt.imshow(image, cmap='gray')
plt.show()
# 마스크를 쓰지 않은 이미지 출력
image = Image.open('./Face-Mask-Classification-20000-Dataset/train/without_mask/seed0030.png')
print("이미지 크기:", image.size)
plt.imshow(image, cmap='gray')
plt.show()
데이터셋 불러오기 transforms_train = transforms.Compose([
transforms.Resize((64, 64)),
transforms.Grayscale(1),
transforms.ToTensor(),
transforms.Normalize((0.5), (0.5)) # normalization
])
data_dir = './Face-Mask-Classification-20000-Dataset/'
train_dataset = datasets.ImageFolder(os.path.join(data_dir, 'train'), transforms_train)
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)
print('Train dataset size:', len(train_dataset))
class_names = train_dataset.classes
print('Class names:', class_names)
기본적인 GAN 실습 - 모델 정의
생성자(Generator) 및 판별자(Discriminator) 모델을 정의합니다.
생성자 클래스: SimpleGenerator
판별자 클래스: SimpleDiscriminator
본 소스코드는 원본 논문에서 제안한 아키텍처와 유사한 아키텍처를 이용합니다.
Reference : PyTorch GAN 예시 코드
latent_dim = 100
n_classes = 2
# 생성자(Generator) 클래스 정의
class SimpleGenerator(nn.Module):
def __init__(self):
super(SimpleGenerator, self).__init__()
# 레이블 정보를 임베딩 차원으로 매핑하여 사용
self.label_embed = nn.Embedding(n_classes, n_classes)
# 하나의 블록(block) 정의
def block(input_dim, output_dim, normalize=True):
layers = [nn.Linear(input_dim, output_dim)]
if normalize:
# 배치 정규화(batch normalization) 수행(차원 동일)
layers.append(nn.BatchNorm1d(output_dim, 0.8))
layers.append(nn.LeakyReLU(0.2, inplace=True))
return layers
# 생성자 모델은 연속적인 여러 개의 블록을 가짐
self.model = nn.Sequential(
*block(latent_dim + n_classes, 128, normalize=False),
*block(128, 256),
*block(256, 512),
*block(512, 1024),
*block(1024, 2048),
nn.Linear(2048, 1 * 64 * 64),
nn.Tanh()
)
def forward(self, noise, labels):
# 노이즈(noise) 벡터와 레이블 임베딩을 순차적으로 연결하여 입력
inputs = torch.cat((noise, self.label_embed(labels)), -1)
img = self.model(inputs)
img = img.view(img.size(0), 1, 64, 64)
return img
# 판별자(Discriminator) 클래스 정의
class SimpleDiscriminator(nn.Module):
def __init__(self):
super(SimpleDiscriminator, self).__init__()
# 레이블 정보를 임베딩 차원으로 매핑하여 사용
self.label_embed = nn.Embedding(n_classes, n_classes)
self.model = nn.Sequential(
nn.Linear(1 * 64 * 64 + n_classes, 512),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(512, 512),
nn.Dropout(0.4),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(512, 512),
nn.Dropout(0.4),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(512, 512),
nn.Dropout(0.4),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(512, 1),
)
# 이미지에 대한 판별 결과를 반환
def forward(self, img, labels):
flattened = img.view(img.size(0), -1)
# 이미지 벡터와 레이블 임베딩을 순차적으로 연결하여 입력
inputs = torch.cat((flattened, self.label_embed(labels)), -1)
output = self.model(inputs)
return output
기본적인 GAN 실습 - 모델 학습 및 샘플링
학습을 위해 생성자와 판별자 모델을 초기화합니다.
적절한 하이퍼 파라미터를 설정합니다.
# 생성자(generator)와 판별자(discriminator) 초기화
simple_generator = SimpleGenerator()
simple_discriminator = SimpleDiscriminator()
simple_generator.cuda()
simple_discriminator.cuda()
# 손실 함수(loss function)
adversarial_loss = nn.MSELoss()
adversarial_loss.cuda()
# 학습률(learning rate) 설정
lr = 0.0002
# 생성자와 판별자를 위한 최적화 함수
optimizer_G = torch.optim.Adam(simple_generator.parameters(), lr=lr, betas=(0.5, 0.999))
optimizer_D = torch.optim.Adam(simple_discriminator.parameters(), lr=lr, betas=(0.5, 0.999))
모델을 학습하면서 주기적으로 샘플링하여 결과를 확인할 수 있습니다.
!mkdir -p ./results/simple/
n_epochs = 100 # 학습의 횟수(epoch) 설정
sample_interval = 500 # 몇 번의 배치(batch)마다 결과를 출력할 것인지 설정
start_time = time.time()
for epoch in range(n_epochs):
for i, (imgs, labels) in enumerate(train_dataloader):
# 진짜(real) 이미지와 가짜(fake) 이미지에 대한 정답 레이블 생성
real = torch.cuda.FloatTensor(imgs.shape[0], 1).fill_(1.0) # 진짜(real): 1
fake = torch.cuda.FloatTensor(imgs.shape[0], 1).fill_(0.0) # 가짜(fake): 0
real_imgs = imgs.cuda()
labels = labels.cuda()
""" 생성자(generator)를 학습합니다. """
optimizer_G.zero_grad()
# 랜덤 노이즈(noise) 및 랜덤 레이블(label) 샘플링
z = torch.normal(mean=0, std=1, size=(imgs.shape[0], latent_dim)).cuda()
generated_labels = torch.randint(0, n_classes, (imgs.shape[0],)).cuda()
# 이미지 생성
generated_imgs = simple_generator(z, generated_labels)
# 생성자(generator)의 손실(loss) 값 계산
g_loss = adversarial_loss(simple_discriminator(generated_imgs, generated_labels), real)
# 생성자(generator) 업데이트
g_loss.backward()
optimizer_G.step()
""" 판별자(discriminator)를 학습합니다. """
optimizer_D.zero_grad()
# 판별자(discriminator)의 손실(loss) 값 계산
real_loss = adversarial_loss(simple_discriminator(real_imgs, labels), real)
fake_loss = adversarial_loss(simple_discriminator(generated_imgs.detach(), generated_labels), fake)
d_loss = (real_loss + fake_loss) / 2
# 판별자(discriminator) 업데이트
d_loss.backward()
optimizer_D.step()
done = epoch * len(train_dataloader) + i
if done % sample_interval == 0:
# 클래스당 8개의 이미지를 생성하여 2 X 8 격자 이미지에 출력
z = torch.normal(mean=0, std=1, size=(n_classes * 8, latent_dim)).cuda()
labels = torch.LongTensor([i for i in range(n_classes) for _ in range(8)]).cuda()
generated_imgs = simple_generator(z, labels)
save_image(generated_imgs, f"./results/simple/{done}.png", nrow=8, normalize=True)
# 하나의 epoch이 끝날 때마다 로그(log) 출력
print(f"[Epoch {epoch}/{n_epochs}] [D loss: {d_loss.item():.6f}] [G loss: {g_loss.item():.6f}] [Elapsed time: {time.time() - start_time:.2f}s]")
from IPython.display import Image as Display
Display('./results/simple/28000.png')
