[DL/모델구현] CNN Autoencoder

## 소스 깃헙 : https://github.com/hyesukim1/Gen_AI

 

컨볼루션 오토인코더란?

- 기본적으로 오토인코더의 구조를 가지고 있지만, 데이터의 특징을 추출하기 위해 컨볼루션 신경망(CNN) 구조를 사용합니다. 

 

구조

1. 인코더

- 인코더는 입력이미지를 점점 작은 차원으로 압축하는 역할

- 컨볼루션 레이어와 풀링 레이어를 사용하여 이미지의 중요한 특징들만 추출하고, 이를 잠재 공간에 표현함

 

2. 디코더

- 디코더는 인코더에서 생성된 압축된 데이터를 다시 원래의 이미지 크기로 복원하는 역할

- 일반적으로 transposed convolution 레이어를 사용

 

* upsampling을 사용할 수 있으나 covolution이 학습할 수 있는 파라미터가 있어서

cnn autoencoder는 ConvTranposed2d를 사용합니다.

upsampling은 데이터를 말그대로 키우기만 하는거고 학습을 시키는 파라미터가 존재하지 않습니다.

torch.nn.Upsample(size=None, scale_factor=None, mode='nearest') mode 인자에 따라 어떤 보간법을 쓸지 정해줍니다.

---

데이터

캐글 데이터: https://www.kaggle.com/competitions/image-matching-challenge-2024

Image Matching Challenge 2024 - Hexathlon | Kaggle

 

데이터 폴더 구조

보통 작업할 데이터는 SSD에 옮겨두고 아래의 폴더구조로 진행했습니다.

 

data

-folder1

  -001.png

  -002.png

  -003.png

-folder2

  -001.png

  -002.png

  -003.png

---

CustomDataset

- 파이토치 프레임워크를 사용하여 CutomDataset 클래스를 만들어주었습니다. 

import torch
from torch import nn
from torchvision import transforms
from torch.utils.data import Dataset
from torch.utils.data import DataLoader

import os
import numpy as np
from glob import glob
from PIL import Image
import matplotlib.pyplot as plt


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 CustomDataset(Dataset):

CustomDataset 클래스는 PyTorch의 Dataset 클래스를 상속 받아 사용합니다. 

이를 통해 __len__과 __getitem__을 오버라이딩합니다. 

 

def __init__(self, data_dir, transform=None):

매개변수 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

CustomDataset 클래스가 선언됨과 동시에 data_dir, image_paths, class_names, 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을 거친 이미지를 리턴합니다.

 

---

Autoencoder

- 파이토치 프레임워크를 사용하여 Autoencoder 클래스를 만들어주었습니다. 

class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()

        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.ReLU()
        )

        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1),
            nn.AdaptiveAvgPool2d((1, 4096)),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1),
            nn.AdaptiveAvgPool2d((1, 8192)),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1),
            nn.AdaptiveAvgPool2d((1, 16384)),
            nn.Sigmoid(),
        )
    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return encoded, decoded

Autoencoder 클래스는 파이토치의 nn.Module을 상속받아 정의됩니다. 이를 통해 신경망 모델의 기존적인 구조와 메소드를 사용할 수 있습니다.

 

super(Autoencoder, self).__init__()

상위 클래스인 nn.Module의 생성자를 호출하여 모듈의 초기화를 수행합니다. nn.Module에 정의된 기본 설정을 상속받기위해서 필요합니다.

 

self.encoder = nn.Sequential(
    nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
    nn.ReLU()
)

오토인코더에 encoder에 해당되는 코드입니다. 컨볼루션 레이어와 비선형 활성화 함수 Relu를 적용하여 인풋 이미지 데이터의 차원을 점차 줄여갑니다.

 

self.decoder = nn.Sequential(
    nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1),
    nn.AdaptiveAvgPool2d((1, 4096)),
    nn.ReLU(),
    nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1),
    nn.AdaptiveAvgPool2d((1, 8192)),
    nn.ReLU(),
    nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1),
    nn.AdaptiveAvgPool2d((1, 16384)),
    nn.Sigmoid(),
)

오토인코더에서 decoder에 해당되는 코드입니다. ConvTranspose2d를 사용하여 입력 데이터의 차원을 확장합니다. 맨 마지막 활성화 함수는 Sigmoid를 사용했습니다.

 

(이슈) ConvTranspose2d의 사이즈 미스 이슈

디코딩 과정에서 각각의 인코더, 디코더 레이어에서 출력되는 사이즈가 달랐습니다.

저는 동일한 이미지 사이즈 128*128로 리사이징했기때문에 입력 이미지별로 사이즈가 다르지 않은 상태였어서 출력값을 AdaptiveAvgPooling을 통해서 pooling layer를 추가하여 출력 값을 고정하였습니다. 

 

링크: https://discuss.pytorch.kr/t/conv2d-convtranspose2d/222

Conv2d, ConvTranspose2d의 사이즈 이슈

 

 

def forward(self, x):
    encoded = self.encoder(x)
    decoded = self.decoder(encoded)
    return encoded, decoded

 

모델의 순전파 함수를 정의합니다. 입력 데이터 x를 받아 인코더와 디코더를 거쳐 출력합니다. 

순전파 함수는 인코딩된 데이터(encoded) 디코딩된 데이터(decoded)를 반환합니다.

---

실제로 학습 하기

def train(autoencoder, train_loader):
    autoencoder.train()
    for step, x in enumerate(train_loader):
        train_x = x.view(-1, 1, 128*128).to(DEVICE)
        train_y = x.view(-1, 1, 128*128).to(DEVICE)

        encoded, decoded = autoencoder(train_x)

        loss = criterion(decoded, train_y) # 스텝별로 로스를 구함
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

실제 학습이 진행되는 train함수를 만들어주었습니다. 

인자로는 모델과 배치 데이터 셋을 받습니다. 

 

저는 gpu에서 학습을 진행하기 때문에 DEVICE를 cuda로 설정하고 데이터를 gpu에 올려주었습니다.

 

pytorch는 model.train()을 통해 학습을 시켜줍니다. 배치 사이즈 만큼 데이터를 반복하여 step별 학습을 진행합니다. 

loss는 mse로 설정하였습니다. 

 

EPOCH = 5
BATCH_SIZE = 8
USE_CUDA = torch.cuda.is_available()
DEVICE = torch.device("cuda" if USE_CUDA else "cpu")
print("Using Device:", DEVICE)

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)

autoencoder = Autoencoder().to(DEVICE)
optimizer = torch.optim.Adam(autoencoder.parameters(), lr=0.005)
criterion = nn.MSELoss()

각종 하이퍼 파라미터들을 한번에 설정했습니다. 

데이터 수가 적어서 Epoch가 5이내로 모델이 빠르게 수렴했기 때문에 5로 설정하였고,

배치 사이즈도 8로 설정했습니다. 

 

이미지 전처리가 진행되는 transform도 설정해주었습니다.

옵티마이저는 Adam, 러닝레이트는 0.005로 설정하였고 위에서 말씀드린 것처럼 로스는 MSE로 설정하였습니다.

 

view_data = [data_loader.dataset[i].view(-1, 1, 128*128).to(DEVICE) for i in range(5)]

for epoch in range(1, EPOCH+1):
    train(autoencoder, data_loader.dataset)

    decoded_dt = []
    for i in view_data:
        _, decoded_data = autoencoder(i)
        decoded_dt.append(decoded_data)

    f, a = plt.subplots(2, 5, figsize=(5, 2))
    print("[Epoch {}]".format(epoch))

    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):
        _, decoded_data = autoencoder(i)
        img = np.reshape(decoded_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.show()

모델이 학습되면서 5개의 다른 이미지가 에포크마다 어떻게 만들어지는지 볼 수 있는 소스입니다.

위 이미지는 원본이미지, 아래 이미지는 모델 통과한 이미지입니다.

 

1_EPOCH

5_EPOCH

학습이 진행될수록 아웃풋 이미지가 조금더 선명해지는걸 알 수 있습니다. 

---

전체 소스

import torch
from torch import nn
from torchvision import transforms
from torch.utils.data import Dataset
from torch.utils.data import DataLoader

import os
import numpy as np
from glob import glob
from PIL import Image
import matplotlib.pyplot as plt

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 Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()

        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.ReLU()
        )

        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1),
            nn.AdaptiveAvgPool2d((1, 4096)),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1),
            nn.AdaptiveAvgPool2d((1, 8192)),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1),
            nn.AdaptiveAvgPool2d((1, 16384)),
            nn.Sigmoid(),
        )
    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return encoded, decoded


def train(autoencoder, train_loader):
    autoencoder.train()
    for step, x in enumerate(train_loader):
        train_x = x.view(-1, 1, 128*128).to(DEVICE)
        train_y = x.view(-1, 1, 128*128).to(DEVICE)

        encoded, decoded = autoencoder(train_x)

        loss = criterion(decoded, train_y) # 스텝별로 로스를 구함
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print('loss:', loss)

EPOCH = 5
BATCH_SIZE = 8
USE_CUDA = torch.cuda.is_available()
DEVICE = torch.device("cuda" if USE_CUDA else "cpu")
print("Using Device:", DEVICE)

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)

autoencoder = Autoencoder().to(DEVICE)
optimizer = torch.optim.Adam(autoencoder.parameters(), lr=0.005)
criterion = nn.MSELoss()

view_data = [data_loader.dataset[i].view(-1, 1, 128*128).to(DEVICE) for i in range(5)]

for epoch in range(1, EPOCH+1):
    train(autoencoder, data_loader.dataset)

    decoded_dt = []
    for i in view_data:
        _, decoded_data = autoencoder(i)
        decoded_dt.append(decoded_data)

    f, a = plt.subplots(2, 5, figsize=(5, 2))
    print("[Epoch {}]".format(epoch))

    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):
        _, decoded_data = autoencoder(i)
        img = np.reshape(decoded_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.show()