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에서 볼 수 있다.