## 소스 깃헙 : https://github.com/hyesukim1/Gen_AI
VAE
- VAE는 오토인코더와 유사한 구조를 가지고 있지만 latent variable 부분에서 잠재 변수에 분포를 정규 분포 형태로 나타내고 그 분포에서 샘플링하도록 합니다.
- VAE는 딥러닝 네트워크로 역전파를 통한 학습이 가능한 아키텍처로 구성합니다. 단순히 정규 분포에서 샘플링하면 인코더와 디코더 사이의 잠재 변수가 미분이 불가능해집니다. 이를 해결하기 위해 인코더 출력부에 μ(평균),σ(분산)를 출력하고 이 값의 선형 결합을 통해 가우시안 분포의 샘플링을 하면 μ, σ에 대해 미분이 가능해집니다.(=Reparametrization Trick)
- 오토인코더는 데이터 압축이 목적이고 VAE는 데이터 생성이 목적입니다.
데이터
지난번과 동일하게 캐글데이터를 사용합니다. 실제 데이터는 RGB의 형태를 가지지만 실제 학습시에는 grayscale로 사용하였습니다.
캐글 데이터: https://www.kaggle.com/competitions/image-matching-challenge-2024
Image Matching Challenge 2024 - Hexathlon | Kaggle
www.kaggle.com
데이터 폴더 구조
data
-folder1
-001.png
-002.png
-003.png
-folder2
-001.png
-002.png
-003.png
데이터 수량/크기
7개 클래스의 2,204장
3.28GB
CustomDataset
- torch.utils.data.Dataset을 상속받아 CutomDataset 클래스를 만들어주었습니다.
커스텀 데이터셋을 만들땐 아래와 같은 구조로 코드를 작성합니다.
class CustomDataset(torch.utils.data.Dataset):
def __init__(self): => 데이터 셋의 전처리를 진행
def __len__(self): => 총 샘플의 수를 적어줌
def __getitem__(self, idx): => 데이터 셋의 실제 샘플을 가져옴
아래의 코드가 실제 제가 만든 CustomDataset 클래스입니다.
class CustomDataset(Dataset):
def __init__(self, data_dir, transform=None):
self.data_dir = data_dir
self.image_paths = glob(os.path.join(data_dir, '*/*.png'))
self.class_names = os.listdir(self.data_dir)
self.transform = transform
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
image_path = self.image_paths[idx]
img = Image.open(image_path).convert('L')
if self.transform:
img = self.transform(img)
return img
transform = transforms.Compose([transforms.Resize((128, 128)),
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor()
])
pth = '이미지_경로'
data = CustomDataset(pth, transform=transform)
data_loader = DataLoader(data, batch_size=BATCH_SIZE, shuffle=True)
def __init__(self, data_dir, transform=None):
self.data_dir = data_dir
self.image_paths = glob(os.path.join(data_dir, '*/*.png'))
self.class_names = os.listdir(self.data_dir)
self.transform = transform
데이터의 경로와 클래스 이름 그리고 어떤 전처리를 할 것인지 정의 해줍니다.
def __getitem__(self, idx):
image_path = self.image_paths[idx]
img = Image.open(image_path).convert('L')
if self.transform:
img = self.transform(img)
return img
주어진 index에 해당하는 샘플을 데이터셋에 불러오고 transform을 거쳐 최종적으로 텐서 형태로 반환해줍니다.
pth = '이미지_경로'
data = CustomDataset(pth, transform=transform)
data_loader = DataLoader(data, batch_size=BATCH_SIZE, shuffle=True)
실제로 사용할 때는 이미지가 있는 경로와 위에서 정의한 transform을 인자로 넣은 후 data라는 변수에 담아줍니다.
DataLoader는 샘플을 미니배치로 전달하거나 에포크마다 데이터를 다시 섞어주거나 Multiprocessing을 사용해 데이터 검색속도를 높여주는 순회 가능한 객체(iterable)입니다.
data, batch_size, shuffle 여부를 DataLoader에 주어 우리가 원하는 조건으로 순회할 수 있도록 만들어줍니다.
VAE
VAE는 인코더, 디코더의 네크워크로 구성되어있어 오토인코더와 비슷한 구조로 되어있습니다.
하지만 바틀넥 부분 즉, 모델의 중간 부분이 다릅니다. 이는 모델의 목적이 다르기 때문인데요.
오토인코더는 주로 차원축소의 목적을 가지고 VAE모델은 확률 분포를 이용해 어떤 새로운 데이터를 만들어내는데 목적이 있습니다.
제가 사용한 캐글 데이터로 모델의 아키텍쳐는 위의 이미지와 같습니다.
latent dimension을 64로 설정하였습니다. autoencoder로 차원축소하여 데이터가 잘 압축되었는지 확인할땐 주로 2차원 또는 3차원까지 줄여서 그래프로 그린다음 확인하는 방식을 사용합니다. 저는 vae로 이미지를 복원하는데 초점을 맞춰 모델을 짜서 너무 낮은 차원으로 줄이진 않았습니다.
그리고, vae가 블러리한 이슈가 있는데 차원을 낮게 가져갈 수록 복원할때 블러리 정도가 심해져서 많이 내리진 않았습니다. 이러한 특징 때문에 vae 이후 모델은 이런 블러 현상을 잡으려고 여러 방법을 시도합니다. (이거는 나중에 정리해보겠습니다.)
class Flatten(nn.Module):
def forward(self, input):
global BATCH_SIZE
return input.view(input.size()[0], -1).to(DEVICE)
class UnFlatten(nn.Module):
def forward(self, input):
global BATCH_SIZE
return input.view(-1, 128, 16, 16).to(DEVICE)
class VAE(nn.Module):
def __init__(self, h_dim=128*16*16, z_dim=64):
super(VAE, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), # 128*16*16
nn.BatchNorm2d(128),
nn.ReLU(),
Flatten()
)
self.fc1 = nn.Linear(h_dim, z_dim).to(DEVICE)
self.fc2 = nn.Linear(h_dim, z_dim).to(DEVICE)
self.fc3 = nn.Linear(z_dim, h_dim).to(DEVICE)
self.decoder = nn.Sequential(
UnFlatten(),
nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(1),
nn.Sigmoid(),
)
def reparameterize(self, mu, logvar):
std = torch.exp(0.5*logvar)
eps = torch.randn_like(std)
return mu + eps*std
def bottleneck(self, h):
mu, logvar = self.fc1(h), self.fc2(h) # be sure not to add activation functions here!
logvar = torch.clamp(logvar, min=-4, max=4)
z = self.reparameterize(mu, logvar)
return z, mu, logvar
def encode(self, x):
h = self.encoder(x)
return self.bottleneck(h)
def decode(self, z):
z = F.relu(self.fc3(z))
z = self.decoder(z)
return z
def forward(self, x):
z, mu, logvar = self.encode(x)
z = self.decode(z)
return z, mu, logvar
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), # 128*16*16
nn.BatchNorm2d(128),
nn.ReLU(),
Flatten()
)
차원을 축소하는 encoder 부분입니다. conv-bn-relu의 순서로 이루어져있고 batchnorm은 학습의 안정화를 위해 추가하였습니다.
batchnorm을 사용할땐 무작위하게 사용하면 안됩니다. 저는 데이터 이미지가 아이코닉한 상태라(이미지 안에 만들려는 객체가 중앙에 위치하여 잘 보임) batchnorm의 효과를 잘보았지만 만약 자신이 사용하는 데이터가 특정 인덱스를 주고 자른 이미지이며 나중에 이미지들을 다 컨켓한다면 batchnorm을 사용해서 학습을 하는것이 의미가 없어질 수 있습니다.
보틀넥 부분에 인풋으로 넣기 위해 Flatten()을 추가하였습니다.
self.fc1 = nn.Linear(h_dim, z_dim).to(DEVICE) # Mu
self.fc2 = nn.Linear(h_dim, z_dim).to(DEVICE) # 로그분산
self.fc3 = nn.Linear(z_dim, h_dim).to(DEVICE) # 디코더 인풋
fc1, fc2는 reparameterize tric을 사용하기 위해 z_dim으로 선형 변환을 시켜줍니다. 평균과 분산을 구한 후 잠재 공간의 샘플을 생성합니다.
def reparameterize(self, mu, logvar):
std = torch.exp(0.5*logvar)
eps = torch.randn_like(std)
return mu + eps*std
eps는 평균이 0이고 표준편차가 1인 정규 분포에서 샘플링 된 값입니다. std와 동일한 모양과 데이터 타입을 가진 텐서를 생성하고 이 텐서의 각 원소는 표준 정규분포에서 무작위로 추출된 값입니다.
실제, 샘플을 생성하는 부분은 리턴 값으로 보내는 mu + eps*std 입니다. 이부분이 입력 데이터에 대응해서 잠재 변수 분포에서 하나의 샘플을 생성하는 것과 같다고 보면 됩니다.
self.decoder = nn.Sequential(
UnFlatten(),
nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(1),
nn.Sigmoid(),
)
선형으로 생성된 샘플을 UnFlatten()시킨 후 16*16*128로 변환해준 후 다시 convtranspose-bn-relu연산으로 샘플을 만들어내는 연산을 진행합니다.
def loss_function(recon_x, x, mu, logvar):
BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return BCE + KLD, BCE, KLDloss function 함수는 reconstruction loss와 KL divergence 두가지 손실 요소를 더해줍니다.
reconstruction error는 VAE가 생성한 출력 이미지와 원본 이미지가 가능한 유사하게 만드는 것입니다.
계산 방식은 이진 크로스 엔트로피로 계산되며 각 픽셀이 독립적인 베르누이 분포를 따른다고 가정할 때 적합합니다.
BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
reduction='sum'으로 준다는 것은 손실 값을 모든 데이터 포인트에 대해 합친다는 것입니다.
BCE를 통해 실제 값과 예측 값의 차이를 구할 수 있습니다.
KL divergence는 VAE의 z_dim이 zero-mean Gaussian이라고 가정했을때 잠재 변수의 사전 분포와 인코더가 제공하는 사후 분포 간의 차이를 측정해 잠재 공간이 더 의미있는 공간으로 가도록 유도합니다.
Train
train_history = {
'total_loss': [],
'bce_loss': [],
'kld_loss': [],
'mu_range_min': [],
'mu_range_max': [],
'logvar_range_min': [],
'logvar_range_max': []
}
for epoch in range(1, EPOCH+1):
model.train()
for step, x in enumerate(data_loader):
images = x.float().to(DEVICE)
recon_images, mu, logvar = model(images)
loss, bce, kld = loss_function(recon_images, images, mu, logvar)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
print(f"Epoch [{epoch + 1}/{EPOCH}], total loss: {loss.item():.4f}, bce: {bce.item():.4f}, kld: {abs(kld.item()):.4f}")
print(f"Mu range: {torch.min(mu[0])} ~ {torch.max(mu[0])}, Logvar range: {torch.min(logvar[0])} ~ {torch.max(logvar[0])}")
if (epoch + 1) % 5 == 0:
torch.save(model.state_dict(), 'D:/result/kaggle/model/'+f'Epoch_{epoch}.h5')
train_history['total_loss'].append(loss.item())
train_history['bce_loss'].append(bce.item())
train_history['kld_loss'].append(abs(kld.item()))
train_history['mu_range_min'].append(torch.min(mu[0]).item())
train_history['mu_range_max'].append(torch.max(mu[0]).item())
train_history['logvar_range_min'].append(torch.min(logvar[0]).item())
train_history['logvar_range_max'].append(torch.max(logvar[0]).item())
np.savetxt('D:/result/kaggle/history/'+f'train_history_epoch_{epoch + 1}.csv',
np.column_stack((train_history['total_loss'],
train_history['bce_loss'],
train_history['kld_loss'],
train_history['mu_range_min'],
train_history['mu_range_max'],
train_history['logvar_range_min'],
train_history['logvar_range_max'])),
delimiter=',',
header='total_loss,bce_loss,kld_loss,mu_range_min,mu_range_max,logvar_range_min,logvar_range_max',
comments='')
f, a = plt.subplots(2, 5, figsize=(5, 2))
for i in range(5):
img = np.reshape(view_data[i].to("cpu").data.numpy(), (128, 128))
a[0][i].imshow(img, cmap='gray')
a[0][i].set_xticks(()); a[0][i].set_yticks(())
for ind, i in enumerate(view_data):
data, _, _ = model(i)
img = np.reshape(data.to("cpu").data.numpy(), (128, 128))
a[1][ind].imshow(img, cmap='gray')
a[1][ind].set_xticks(()); a[1][ind].set_yticks(())
plt.savefig('경로/'+f'Epoch_{epoch}.png', bbox_inches='tight', pad_inches=0.1)실제 학습 소스이고 5에포크마다 결과 이미지, 모델, 히스토리를 저장했습니다.
에포크는 500으로 진행했으나 300이후로 수렴이 안되서 300부근에서 학습을 중단했습니다.
처음엔 큰 특징만 잡아내다가 시간이 지날수록 객체와 배경을 점점 구현하는 모습이 보입니다.
학습 마지막 부분엔 2, 4번째 이미지가 완전히 배경과 건물, 나무가 완전히 섞여버립니다. 이때 쯤 학습을 중단했습니다.
전체 학습 소스
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torchvision import transforms
import time
import os
import numpy as np
from glob import glob
from PIL import Image
import matplotlib.pyplot as plt
from tqdm import tqdm
from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau
USE_CUDA = torch.cuda.is_available()
DEVICE = torch.device("cuda:0" if USE_CUDA else "cpu")
print("Using Device:", DEVICE)
class CustomDataset(Dataset):
def __init__(self, data_dir, transform=None):
self.data_dir = data_dir
self.image_paths = glob(os.path.join(data_dir, '*/*.png'))
self.class_names = os.listdir(self.data_dir)
self.transform = transform
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
image_path = self.image_paths[idx]
img = Image.open(image_path).convert('L')
if self.transform:
img = self.transform(img)
return img
class Flatten(nn.Module):
def forward(self, input):
global BATCH_SIZE
return input.view(input.size()[0], -1).to(DEVICE)
class UnFlatten(nn.Module):
def forward(self, input):
global BATCH_SIZE
return input.view(-1, 128, 16, 16).to(DEVICE)
class VAE(nn.Module):
def __init__(self, h_dim=128*16*16, z_dim=64):
super(VAE, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), # 128*16*16
nn.BatchNorm2d(128),
nn.ReLU(),
Flatten()
)
self.fc1 = nn.Linear(h_dim, z_dim).to(DEVICE) # for mu right before reparameterization
self.fc2 = nn.Linear(h_dim, z_dim).to(DEVICE) # for logvar right before reparameterization
self.fc3 = nn.Linear(z_dim, h_dim).to(DEVICE) # right before decoding starts
self.decoder = nn.Sequential(
UnFlatten(),
nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(1),
nn.Sigmoid(),
)
def reparameterize(self, mu, logvar):
std = torch.exp(0.5*logvar)
eps = torch.randn_like(std)
return mu + eps*std
def bottleneck(self, h):
mu, logvar = self.fc1(h), self.fc2(h) # be sure not to add activation functions here!
logvar = torch.clamp(logvar, min=-4, max=4)
z = self.reparameterize(mu, logvar)
return z, mu, logvar
def encode(self, x):
h = self.encoder(x)
return self.bottleneck(h)
def decode(self, z):
z = F.relu(self.fc3(z))
z = self.decoder(z)
return z
def forward(self, x):
z, mu, logvar = self.encode(x)
z = self.decode(z)
return z, mu, logvar
def loss_function(recon_x, x, mu, logvar):
BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return BCE + KLD, BCE, KLD
EPOCH = 500
BATCH_SIZE = 64
transform = transforms.Compose([transforms.Resize((128, 128)),
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor()
])
pth = '경로/'
data = CustomDataset(pth, transform=transform)
data_loader = DataLoader(data, batch_size=BATCH_SIZE, shuffle=True)
model = VAE()
model = model.to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scheduler = StepLR(optimizer, step_size=20, gamma=0.9)
# scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, verbose=True)
view_data = [data_loader.dataset[i].unsqueeze(0).to(DEVICE) for i in range(5)]
train_history = {
'total_loss': [],
'bce_loss': [],
'kld_loss': [],
'mu_range_min': [],
'mu_range_max': [],
'logvar_range_min': [],
'logvar_range_max': []
}
for epoch in range(1, EPOCH+1):
model.train()
for step, x in enumerate(data_loader):
images = x.float().to(DEVICE)
recon_images, mu, logvar = model(images)
loss, bce, kld = loss_function(recon_images, images, mu, logvar)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
print(f"Epoch [{epoch + 1}/{EPOCH}], total loss: {loss.item():.4f}, bce: {bce.item():.4f}, kld: {abs(kld.item()):.4f}")
print(f"Mu range: {torch.min(mu[0])} ~ {torch.max(mu[0])}, Logvar range: {torch.min(logvar[0])} ~ {torch.max(logvar[0])}")
if (epoch + 1) % 5 == 0:
torch.save(model.state_dict(), '경로/'+f'Epoch_{epoch}.h5')
train_history['total_loss'].append(loss.item())
train_history['bce_loss'].append(bce.item())
train_history['kld_loss'].append(abs(kld.item()))
train_history['mu_range_min'].append(torch.min(mu[0]).item())
train_history['mu_range_max'].append(torch.max(mu[0]).item())
train_history['logvar_range_min'].append(torch.min(logvar[0]).item())
train_history['logvar_range_max'].append(torch.max(logvar[0]).item())
np.savetxt('경로/'+f'train_history_epoch_{epoch + 1}.csv',
np.column_stack((train_history['total_loss'],
train_history['bce_loss'],
train_history['kld_loss'],
train_history['mu_range_min'],
train_history['mu_range_max'],
train_history['logvar_range_min'],
train_history['logvar_range_max'])),
delimiter=',',
header='total_loss,bce_loss,kld_loss,mu_range_min,mu_range_max,logvar_range_min,logvar_range_max',
comments='')
f, a = plt.subplots(2, 5, figsize=(5, 2))
for i in range(5):
img = np.reshape(view_data[i].to("cpu").data.numpy(), (128, 128))
a[0][i].imshow(img, cmap='gray')
a[0][i].set_xticks(()); a[0][i].set_yticks(())
for ind, i in enumerate(view_data):
data, _, _ = model(i)
img = np.reshape(data.to("cpu").data.numpy(), (128, 128))
a[1][ind].imshow(img, cmap='gray')
a[1][ind].set_xticks(()); a[1][ind].set_yticks(())
plt.savefig('경로/'+f'Epoch_{epoch}.png', bbox_inches='tight', pad_inches=0.1)'Train 학습할레나' 카테고리의 다른 글
| [DL/모델구현] CNN Autoencoder (0) | 2024.04.11 |
|---|
