Deep Learning

PGGAN, StyleGAN 설명, Torch 코드 구현

Kimhj 2024. 2. 19. 08:52
  • PGGAN(PixelGAN)
    • PGGAN은 "Progressive Growing of GANs"의 약자로, 픽셀 수준의 이미지 생성에 중점을 둔 모델
    • PGGAN은 단계적으로 해상도를 키우면서 이미지를 생성하는 방식이며, 먼저 매우 작은 해상도의 이미지부터 시작하여 점진적으로 해상도를 키워가며 생성함.
    • 각 단계에서는 이미지 생성에 사용되는 신경망의 크기와 깊이를 조절하며 작은 해상도의 이미지에서는 작은 신경망을 사용하고, 점진적으로 더 큰 신경망을 사용하여 더 높은 해상도의 이미지를 생성함.
    • 이러한 접근 방식은 높은 해상도의 이미지를 생성하는 안정적으로 학습할 있음.
  • Torch 코드 구현
import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, latent_dim, channels):
        super(Generator, self).__init__()
        self.latent_dim = latent_dim
        self.channels = channels
        self.main = nn.Sequential(
            nn.ConvTranspose2d(latent_dim, 512, 4, 1, 0),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.ConvTranspose2d(512, 256, 4, 2, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.ConvTranspose2d(256, 128, 4, 2, 1),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.ConvTranspose2d(128, channels, 4, 2, 1),
            nn.Tanh()
        )

    def forward(self, input):
        return self.main(input)

class Discriminator(nn.Module):
    def __init__(self, channels):
        super(Discriminator, self).__init__()
        self.channels = channels
        self.main = nn.Sequential(
            nn.Conv2d(channels, 128, 4, 2, 1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, 4, 2, 1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 512, 4, 2, 1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(512, 1, 4, 1, 0),
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input)

# Example Usage
latent_dim = 100
channels = 3
generator = Generator(latent_dim, channels)
discriminator = Discriminator(channels)

# Random Noise
noise = torch.randn(1, latent_dim, 1, 1)
fake_image = generator(noise)

# output
output = discriminator(fake_image)

 

 

PGGAN Architecure

 

 

  • StyleGAN
    • StyleGAN은 GAN 기반의 이미지 생성 모델로, 높은 해상도의 이미지를 생성하는 데 중점을 둠.
    • StyleGAN은 이미지 생성에 스타일과 구조를 분리하여 제어하는 방식을 도입했는데, 이를 통해 더 현실적이고 다양한 이미지를 생성할 수 있음.
    • 스타일 전이(Style Transfer) 개념을 활용하여, 이미지의 스타일 정보를 입력으로 제공하고, 이를 기반으로 생성된 이미지의 스타일을 조정할 수 있음.
    • 랜덤 벡터를 입력으로 받는 대신, 스타일 및 잡음 벡터로부터 이미지를 생성하는데 사용되는 각 레이어의 가중치를 조정하여 이미지 생성의 품질을 향상
    • StyleGAN2에서는 progressive growing 유사한 개념을 도입하여 해상도를 점진적으로 키우는 방식을 채택하여 효율성과 안정성을 향상시켰으며, 현재는 v3 까지 나와있음.

 

  • Torch 코드 구현
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

# Pixel normalization layer
class PixelNorm(nn.Module):
    def __init__(self):
        super(PixelNorm, self).__init__()

    def forward(self, x):
        return x / torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + 1e-8)

# Generator block
class GeneratorBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(GeneratorBlock, self).__init__()
        self.conv1 = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.pixel_norm = PixelNorm()
        self.leaky_relu = nn.LeakyReLU(0.2)

    def forward(self, x):
        x = self.pixel_norm(x)
        x = self.leaky_relu(self.conv1(x))
        x = self.pixel_norm(x)
        x = self.leaky_relu(self.conv2(x))
        return x

# Generator
class Generator(nn.Module):
    def __init__(self, latent_dim, channels, base_channels, num_blocks):
        super(Generator, self).__init__()
        self.num_blocks = num_blocks
        self.initial_block = nn.ConvTranspose2d(latent_dim, base_channels * 4, kernel_size=4, stride=4)
        self.blocks = nn.ModuleList()
        for _ in range(num_blocks):
            self.blocks.append(GeneratorBlock(base_channels * 4, base_channels * 4))
        self.final_conv = nn.Conv2d(base_channels * 4, channels, kernel_size=1)

    def forward(self, x):
        x = self.initial_block(x)
        for block in self.blocks:
            x = block(x)
        x = self.final_conv(x)
        return torch.tanh(x)

# Example usage
latent_dim = 100
channels = 3
base_channels = 32
num_blocks = 4

generator = Generator(latent_dim, channels, base_channels, num_blocks)
latent_vector = torch.randn(1, latent_dim, 1, 1)
fake_image = generator(latent_vector)

 

 

StyleGAN v2 Architecture

 

'Deep Learning' 카테고리의 다른 글

Pytorch Pix2Pix 모델 구현  (0) 2024.02.16
[TORCH_UECUDA_DSA] CUDA 오류  (0) 2024.02.05
conda env 실행 시 오류 해결  (0) 2024.01.31
pytorch gpu(cuda) 정보 확인  (0) 2024.01.17
anaconda 다중 사용자(multi-user) 환경 세팅  (0) 2024.01.15