TF2와 Pytorch 공부를 위해 Convolutional Autoencoder를 두 라이브러리 모두에서 작성했습니다.

소개

Autoencoder는 특징량 추출이나 이상 검출 등에 사용되는 신경망의 모델
큰 특징으로서 입력과 출력의 형태가 같고, 그것보다 낮은 차원의 중간층을 포함하고 있다
입력과 출력이 동일하도록 학습시킴으로써 중간층에서 낮은 차원에서 이미지의 특징을 나타낼 수 있게 된다
또, 학습한 데이터와 크게 다른 데이터는 잘 출력으로 재구성할 수 없기 때문에, 파손 등의 이상 검지에도 사용할 수 있다(같다)



일반 Autoencoder는 기본적으로 모든 층이 전체 결합층을 이용하고 있지만, 이번에 사용한 Convolutional Autoencoder는 그 이름대로 컨볼루션층에 의해 차원을 줄이고 있다

모델 구성

이번은 컨벌루션층 3층으로 작게 하고 나서의 전체 결합층, 그리고 복원에 전치 컨벌루션층 3층의 구성으로 모델을 짜

재구성 화상의 깨끗함을 중시한다면, batch normalization 넣거나 중간층의 전체 결합층을 벗겨내는 것이 loss는 낮아지는 것 같지만, 중간층의 차원을 낮게 하면서 재구성 화상도 적당히 하는 것을 목표로 한다 그래서 이런 모양이 됐어

TF2와 Pytorch에 있어서 사이즈등 같게 되도록 파라미터를 조정했지만, 정확히 같은 파라미터에서는 미묘하게 중간층의 사이즈등이 바뀌어 버렸기 때문에 pytorch의 전치 컨벌루션은 padding과 output_padding으로 사이즈를 맞추게 되었다
이 근처의 처리로 미묘하게 변화가 있을지도 모르지만, 지금 하나 이해 부족이므로 유식자 쪽이 있으면 가르쳐 주셨으면 싶다

데이터 세트

학습에는 cifar10을 사용했습니다.
어느 라이브러리에도 자동으로 다운로드 해 주는 기구가 있으므로 데이터의 준비는 비교적 간단
주의점은, 처음에는 다운로드가 필요한 것, Pytorch 쪽은 지금 있는 장소에 디렉토리를 만들어 거기에 데이터를 다운로드해 전개하는 것

TF2

import tensorflow as tf
(train_data, train_label), (test_data, test_label) = tf.keras.datasets.cifar10.load_data()

Pytorch

import torchvision
import torchvision.transforms as transforms
transform = torchvision.transforms.ToTensor()
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)

tensorflow2에서 구현

tf.keras.Sequential에서 encoder와 decoder를 각각 구성하여 call에서 출력(생성 이미지)과 중간층의 값을 반환하도록 했다.

class CAE(Model):
    def __init__(self):
        super(CAE,self).__init__()
        self.encoder = tf.keras.Sequential([
            Conv2D(16, 3, strides=(2,2), padding='same', activation=tf.nn.tanh), # 16x16x16
            Conv2D(32, 3, strides=(2,2), padding='same', activation=tf.nn.tanh), # 8x8x32
            Conv2D(32, 3, strides=(2,2), padding='same', activation=tf.nn.tanh), # 4x4x32
            Flatten(),
            Dense(4*4*8, activation=tf.nn.tanh), # 128 dim.
        ])
        self.decoder = tf.keras.Sequential([
            Dense(4*4*32, activation=tf.nn.tanh),
            Reshape((4,4,32)),
            Conv2DTranspose(32, 3, strides=2, padding='same', activation=tf.nn.tanh), # 8x8x32
            Conv2DTranspose(16, 3, strides=2, padding='same', activation=tf.nn.tanh), # 16x16x16
            Conv2DTranspose(3, 3, strides=2, padding='same', activation=tf.nn.sigmoid), # 32x32x3
        ])

    def call(self, x):
        z = self.encoder(x)
        x_pred = self.decoder(z)
        return x_pred, z

원본 이미지와 생성된 이미지

Pytorch에서 구현

nn.Sequential에서 encoder와 decoder를 각각 구성해, forward로 출력(생성 화상)과 중간층의 값을 돌려주도록 했다

class CAE(nn.Module):
    def __init__(self):
        super(CAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1), nn.Tanh(),
            nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1), nn.Tanh(),
            nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1), nn.Tanh(),
            nn.Flatten(),
            nn.Linear(4*4*32, 128), nn.Tanh(),
        )

        self.decoder = nn.Sequential(
            nn.Linear(128, 4*4*32), nn.Tanh(),
            Reshape(-1, 32, 4, 4),
            nn.ConvTranspose2d(32, 32, kernel_size=3, stride=2, padding=1, output_padding=1), nn.Tanh(),
            nn.ConvTranspose2d(32, 16, kernel_size=3, stride=2, padding=1, output_padding=1), nn.Tanh(),
            nn.ConvTranspose2d(16, 3, kernel_size=3, stride=2, padding=1, output_padding=1), nn.Sigmoid(),
        )

    def forward(self, x):
        z = self.encoder(x)
        x_pred = self.decoder(z)
        return x_pred, z

원본 이미지와 생성된 이미지

감상

전체적으로 비슷하기 때문에, 미묘하게 다른 점에서 실수 실수를 하자 (Sequential() 안에 []가 필요한지 등)
이 정도면 모델보다 학습 돌리거나 데이터 준비하거나하는 코드를 쓰는 것이 더 어렵습니다.

코드 전문은 github 의 코드를 참조