GANs (Generative Adversarial Networks) 예제 (251125)

파일 구조

generator.py
discriminator.py
sampler.py
train_gan.py (학습 + 체크포인트 저장 포함)
load_models.py (저장된 체크포인트 복원 및 샘플 생성)

생성자 : generator.py

import torch
import torch.nn as nn


class Generator(nn.Module):
	"""간단한 Fully-connected Generator 예제

	아키텍처 설명 (초보자 친화적):
	1) 입력: latent_dim 크기의 벡터 z (무작위 노이즈)
	2) 여러 개의 선형 계층(Linear)과 활성화 함수(LeakyReLU)로 신호를 증가시킴
	3) BatchNorm을 사용해 학습 안정화
	4) 마지막 레이어에서 이미지 크기 (1*28*28)로 매핑하고 tanh 활성화를 사용


	왜 tanh를 쓰는가?
	- MNIST 이미지를 -1..1 범위로 정규화했기 때문에 출력도 같은 범위가 되도록 tanh를 사용했습니다.
"""


	def __init__(self, latent_dim: int = 100, img_shape=(1, 28, 28)):
		super().__init__()
		self.latent_dim = latent_dim
		self.img_shape = img_shape


		# nn.Sequential을 사용해 레이어를 쭉 나열합니다. 구조가 간단할 때 유용합니다.
		self.model = nn.Sequential(
		# 처음에는 latent_dim을 작은 은닉 유닛으로 확장
		nn.Linear(latent_dim, 128),
		nn.LeakyReLU(0.2, inplace=True),


		# 더 큰 표현을 얻기 위해 차원을 점진적으로 늘립니다.
		nn.Linear(128, 256),
		nn.BatchNorm1d(256),
		nn.LeakyReLU(0.2, inplace=True),


		nn.Linear(256, 512),
		nn.BatchNorm1d(512),
		nn.LeakyReLU(0.2, inplace=True),


		# 마지막으로 이미지 전체 픽셀 수로 매핑합니다.
		nn.Linear(512, int(torch.prod(torch.tensor(img_shape)))),
		nn.Tanh(), # 출력 범위를 -1..1로 제한
		)


	def forward(self, z: torch.Tensor) -> torch.Tensor:
		"""생성자 전달 함수


		Args:
		z: (batch_size, latent_dim) 형태의 노이즈 텐서
		Returns:
		img: (batch_size, 1, 28, 28) 형태의 생성된 이미지 텐서
		"""
		img = self.model(z)
		# view로 텐서를 이미지 형태로 reshape
		return img.view(img.size(0), *self.img_shape)

판별자 : discriminator.py

"""
Discriminator 모델 정의 파일
- 입력: 이미지 (실제 또는 생성된 이미지)
- 출력: 0~1 확률 (이미지가 진짜일 확률)


간단한 Fully-connected Discriminator를 사용합니다. 실제 환경에서는 CNN(예: DCGAN)을 선호합니다.
"""


import torch
import torch.nn as nn


class Discriminator(nn.Module):
	"""간단한 Fully-connected Discriminator


	아키텍처 설명:
	1) 이미지를 일렬로(flatten) 펴서 Linear 계층에 연결
	2) LeakyReLU로 비선형성 추가
	3) 마지막에 Sigmoid로 확률 출력
	"""


	def __init__(self, img_shape=(1, 28, 28)):
		super().__init__()
		self.img_shape = img_shape
		self.input_dim = int(torch.prod(torch.tensor(img_shape)))


		self.model = nn.Sequential(
			nn.Linear(self.input_dim, 512),
			nn.LeakyReLU(0.2, inplace=True),


			nn.Linear(512, 256),
			nn.LeakyReLU(0.2, inplace=True),


			nn.Linear(256, 1),
			nn.Sigmoid(), # 0~1 사이 확률 출력
			)


	def forward(self, img: torch.Tensor) -> torch.Tensor:
		"""판별자 전달함수


		Args:
		img: (batch_size, 1, 28, 28) 형태의 이미지 텐서
		Returns:
		output: (batch_size, 1) 형태의 확률 텐서
		"""
		img_flat = img.view(img.size(0), -1) # flatten
		return self.model(img_flat)

샘플 시각화 : sampler.py

"""
이미지 샘플링 및 표시를 담당하는 클래스
- 학습 중간에 생성 이미지를 시각화하거나, 학습 후 모델을 불러와 이미지를 생성할 때 사용합니다.
"""

import torch
import matplotlib.pyplot as plt


class Sampler:
	"""Generator로부터 이미지를 생성하고 matplotlib로 출력하는 유틸리티 클래스

	사용법:
	sampler = Sampler(generator, latent_dim=100, 	device=device)
	sampler.sample_images(n=25, filename=None) # 화면에 표시
	sampler.sample_images(n=16, filename='samples.png') # 파일로 저장
	"""


def __init__(self, generator, latent_dim: int = 100, device: str = "cpu"):
self.generator = generator
self.latent_dim = latent_dim
self.device = device


	def sample_images(self, n: int = 25, filename: str = None):
	"""n개의 이미지를 생성하여 5x5 그리드로 표시합니다.


	Args:
	n: 생성할 이미지 수 (기본 25)
	filename: None이면 화면에 표시, 문자열이면 해당 파일로 저장
	"""
		# generator를 평가 모드로 바꿔 Dropout/BatchNorm 등 동작을 추론용으로 변경
		self.generator.eval()


		# 무작위 노이즈 생성
		z = torch.randn(n, self.latent_dim).to(self.device)
		with torch.no_grad():
			gen_imgs = self.generator(z).cpu()


		# 이미지 그리기
		grid_size = int(n ** 0.5)
		plt.figure(figsize=(grid_size, grid_size))
		for i in range(n):
			plt.subplot(grid_size, grid_size, i + 1)
			plt.imshow(gen_imgs[i].squeeze(), cmap='gray', vmin=-1, vmax=1)
			plt.axis('off')
			plt.tight_layout()


		if filename:
			plt.savefig(filename)
			print(f"Saved sample grid to {filename}")
		else:
			plt.show()

# 다시 학습 모드로 돌려놓지 않습니다. 호출자가 필요하면 다시 설정하세요.

GANs 학습 : train_gan.py

"""
fake = torch.zeros(imgs.size(0), 1, device=device)


# ------------------
# 1) Generator 업데이트
# ------------------
optimizer_G.zero_grad()


# 무작위 노이즈 생성
z = torch.randn(imgs.size(0), latent_dim, device=device)
# Generator로 가짜 이미지 생성
gen_imgs = generator(z)
# 판별자가 가짜 이미지를 진짜로 판단하도록 만들기 위해 valid 레이블 사용
g_loss = adversarial_loss(discriminator(gen_imgs), valid)
# 역전파 및 파라미터 갱신
g_loss.backward()
optimizer_G.step()


# ------------------
# 2) Discriminator 업데이트
# ------------------
optimizer_D.zero_grad()


# 실제 이미지에 대한 판별 손실
real_loss = adversarial_loss(discriminator(real_imgs), valid)
# 가짜 이미지에 대한 판별 손실 (detach로 generator의 그래프와 분리)
fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)
d_loss = (real_loss + fake_loss) / 2


d_loss.backward()
optimizer_D.step()


# epoch 끝나면 로그 출력
print(f"Epoch {epoch}/{epochs} | D_loss: {d_loss.item():.4f} | G_loss: {g_loss.item():.4f}")


# 매 5 epoch마다 샘플 이미지 저장
if epoch % 5 == 0:
	sample_path = os.path.join('samples', f'epoch_{epoch}.png')
	os.makedirs('samples', exist_ok=True)
	sampler.sample_images(n=25, filename=sample_path)


# 매 5 epoch마다 모델 체크포인트 저장
if epoch % 5 == 0:
	ckpt_path = os.path.join(checkpoint_dir, f'gan_epoch_{epoch}.pth')
	save_checkpoint(generator, discriminator, optimizer_G, optimizer_D, epoch, ckpt_path)


# 학습 종료 후 마지막 모델 저장
final_ckpt = os.path.join(checkpoint_dir, 'gan_final.pth')
save_checkpoint(generator, discriminator, optimizer_G, optimizer_D, epochs, final_ckpt)
print('Training finished.')


if __name__ == '__main__':
	device = 'cuda' if torch.cuda.is_available() else 'cpu'
	train_gan(epochs=20, batch_size=64, latent_dim=100, lr=0.0002, device=device)

저장된 모델을 불러와서 활용하기 : load_model.py

"""
저장된 체크포인트를 불러와 Generator와 Discriminator를 복원하는 스크립트
- Generator와 Discriminator의 state_dict를 로드
- 샘플 이미지를 생성하여 저장


사용법 예시:
python load_models.py --checkpoint ./checkpoints/gan_epoch_20.pth --out samples/loaded.png


주의: 이 스크립트는 generator/discriminator 클래스 정의가 동일할 때만 정상 동작합니다.
"""


import argparse
import torch


from generator import Generator
from discriminator import Discriminator
from sampler import Sampler




def load_checkpoint(path: str, device: str = 'cpu'):
	"""체크포인트 로드


	Returns:
	dict with keys: epoch, generator_state_dict, discriminator_state_dict, ...
	"""
	checkpoint = torch.load(path, map_location=device)
	return checkpoint




def restore_models(checkpoint_path: str, latent_dim: int = 100, device: str = 'cpu'):
	"""체크포인트로부터 모델을 복원하고 샘플을 생성합니다."""
	ckpt = load_checkpoint(checkpoint_path, device)


	# 모델 인스턴스 생성
	generator = Generator(latent_dim=latent_dim).to(device)
	discriminator = Discriminator().to(device)


	# state_dict 로드
	generator.load_state_dict(ckpt['generator_state_dict'])
	discriminator.load_state_dict(ckpt['discriminator_state_dict'])


	# 옵티마이저나 epoch 정보가 필요하면 ckpt의 필드를 사용 가능
	epoch = ckpt.get('epoch', None)
	print(f"Loaded checkpoint from {checkpoint_path}. epoch={epoch}")


	return generator, discriminator, epoch




if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument('--checkpoint', type=str, required=True, help='path to checkpoint .pth file')
	parser.add_argument('--latent_dim', type=int, default=100, help='latent dim used in generator')
	parser.add_argument('--out', type=str, default=None, help='output image filename to save samples')
	args = parser.parse_args()


	device = 'cuda' if torch.cuda.is_available() else 'cpu'
	gen, disc, epoch = restore_models(args.checkpoint, latent_dim=args.latent_dim, device=device)


	sampler = Sampler(gen, latent_dim=args.latent_dim, device=device)
	# 기본적으로 16개 이미지를 생성
	out = args.out if args.out else f'loaded_samples_epoch_{epoch}.png'
	sampler.sample_images(n=16, filename=out)

GANs 생성기는...

GAN 생성기는 특정 숫자(0~9)를 선택해서 생성하는 것이 아니라, MNIST 데이터 분포 전체를 따라 “숫자처럼 보이는 임의의 이미지”를 생성한다. 즉, 조건 없이(Condition 없이) 랜덤 노이즈만 입력받아 아무 숫자 형태나 생성하게 된다.

생성기(Generator)의 입력은 다음과 같다.

z = torch.randn(batch_size, latent_dim)

z는 랜덤 노이즈이며, 레이블(label)에 대한 정보가 전혀 없다.

따라서 모델은 “이 노이즈를 MNIST 분포에 가장 비슷한 이미지로 바꾸는 법”만 학습한다.
MNIST 분포 = 0,1,2,...,9 모든 숫자가 등장하는 분포,
→ 그래서 생성기는 0~9 중 어떤 형태라도 랜덤하게 생성하게 된다.

판별자는 숫자 종류를 구분하지 않는다.
판별자는 클래스를 구분하는 분류기(classifier)가 아니라
실제 데이터 분포의 특징을 학습하는 binary classifier 이다.

판별자가 보는 특징 예시:
MNIST 숫자는 가운데 부분이 검정색일 가능성이 높다.
바깥쪽은 배경이라 흰색에 가깝다.
획 두께, 곡선 패턴, 숫자 구조 등이 MNIST에서 자주 등장한다.

이런 특징을 학습할 뿐
“이건 2다 / 9다”는 전혀 모른다.

특정 숫자를 생성하려면? : Conditional GAN(cGAN)

입력 = (노이즈 z + 원하는 숫자 레이블 y)
출력 = y에 해당하는 숫자 이미지

예: “7을 생성하고 싶다”
→ z + "7" 레이블을 입력으로 넣어 “7” 이미지 생성.