AutoEncoder: 이미지 압축 및 노이즈 제거

AutoEncoder

AutoEncoder는 Encoder를 통해 정보를 압축하고, Decoder를 통해 정보를 복원한다. 예를 들어, 28 x 28 흑백 이미지가 있다고 해보자. 이미지는 총 784개의 정보를 가진다. 만약 이미지를 64개의 정보로 압축할 수 있다면 어떨까?

Encoder는 Linear layer를 이용해 784개의 정보를 64개로 압축한다. 압축된 정보는 latent variable라고 하며, $z$로 표기한다.

nn.Linear(784, 64)

Decoder는 64개로 압축된 latent variable을 다시 원본인 784개 정보로 복원한다.

nn.Linear(64, 784)

 

왜 사용할까?

정보를 다시 원본으로 복원시킬 거라면 왜 사용할까? Decoder가 원본보다 더 나은 형태로 복원시킬 수 있기 때문이다. 예를 들어, 노이즈가 있는 이미지를 노이즈가 없는 이미지로 복원할 수 있다. 또는 워터마크가 있는 이미지를 워터마크 없는 이미지로 복원할 수 있다. 이 개념은 발전되어 이후 생성형 딥러닝 모델으로 확장된다. 

본 글에서는 노이즈가 낀 MNIST 이미지를 받아, 노이즈를 제거하는 모델을 학습한다.


모델 학습 및 분석

  • 모델: Encoder + Decoder로 구성된 AutoEncoder
  • 데이터: MNIST 숫자 손글씨 이미지
  • 입력: 노이즈가 추가된 1D 이미지 (28 * 28,)
  • 출력: 생성한 1D 이미지 (28 * 28,)
  • 정답: 노이즈가 없는 1D 이미지 (28 * 28,)

데이터 전처리

def add_noise(inputs):
    noise = torch.randn_like(inputs) * noise_factor
    return inputs + noise
    
def preprocess(inputs):
    inputs = inputs.view(inputs.size(0), -1)
    noisy_inputs = add_noise(inputs)
    noisy_inputs = torch.clamp(noisy_inputs, -1.0, 1.0)
    return inputs, noisy_inputs

이미지(inputs)에 노이즈를 추가한다. 그리고 Linear layer에 입력으로 주기 위해 (28, 28) 형태의 2D 이미지를 (784,) 형태로 변환한다. 입력 이미지로 활용할 noisy_inputs는 [-1, 1] 범위로 정규화한다. 이미지 밝기에 따라 분포가 다를 수 있기 때문에 일정한 분포를 가지도록 정규화하는 과정은 성능 향상에 도움을 준다.

모델 구조

class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28 * 28, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
        )
        self.decoder = nn.Sequential(
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 28 * 28),
            nn.Tanh(),
        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded
  • Encoder: 이미지를 (28 * 28) →  256 → 128 → 64로 압축
  • Decoder: latent variable을 64   128 → 256 → 28 * 28로 복원

앞서 설명했듯, Linear layer를 이용해 정보를 압축/복원하는 아주 간단한 구조다.

학습

  • learning rate: 0.001
  • epoch: 10
  • batch: 128
  • loss: MSE
  • optimizer: Adam
model = Autoencoder()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(n_epochs):
    model.train()
    for images, _ in train_loader:
        real_images, noisy_images = preprocess(images)
        real_images = real_images.to(device)
        noisy_images = noisy_images.to(device)

        outputs = model(noisy_images)
        loss = criterion(outputs, real_images)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

결과

real_images = real_images.cpu().view(-1, 28, 28)
noisy_images = noisy_images.cpu().view(-1, 28, 28)
outputs = outputs.cpu().view(-1, 28, 28)

직관적인 이해를 위해 (28 * 28,)의 1D 이미지를 (28, 28)의 2D 이미지로 시각화했다.

출력은 입력에서 노이즈가 제거된 모습이다. 원본과 유사하지만 약간 블러한 듯한 모습을 띤다.


전체 코드는 Github: autoencoder_mnist.ipynb에서 볼 수 있다.