오토인코더는 라벨 없이(비지도 방법, unsupervised) 원하는 출력값(reconstructed image)을 생성하는 모델임. (아래는 MNIST 데이터를 사용한 autoencoder!)
1. 라이브러리
import torch
import torchvision
from torchvision import transforms
import torch.nn.functional as F
import torch.nn as nn
import torch.optim as optim
import numpy as np
import cv2
import matplotlib.pyplot as plt
# CPU/GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f'{device} is available.')2. 데이터(MNIST) 불러오기
# MNIST는 흑백 이미지로 (1x28x28)의 형태
dataset = torchvision.datasets.MNIST('./data/', download=True, train=True, transform=transforms.ToTensor())
trainloader = torch.utils.data.DataLoader(dataset,batch_size=50,shuffle=True)
print(dataset[0][0].shape) # 이미지 하나의 shape 찍어보면
#result
torch.Size([1, 28, 28])3. 모델 생성 (Flatten과 Deflatten 만들고 -> Autoencoder)
Flatten, Deflatten 들어가는 위치 파악 !
class Flatten(torch.nn.Module): # 4D -> 2D로 flatten
def forward(self, x):
batch_size = x.shape[0]
return x.view(batch_size, -1)
# (배치 수, 채널 수, 이미지 너비, 이미지 높이) ->
# (배치 수, 채널 수 * 이미지 너비 * 이미지 높이)
class Deflatten(nn.Module): # 2D -> 4D로 deflatten
def __init__(self, k):
super(Deflatten, self).__init__()
self.k = k
def forward(self, x):
s = x.size()
# 벡터 사이즈 = 채널 수*이미지 너비*이미지 높이
# 벡터 사이즈 = 채널 수*이미지 사이즈**2
# 이미지 사이즈 = (벡터 사이즈//채널 수)**.5
feature_size = int((s[1]//self.k)**.5)
return x.view(s[0], self.k, feature_size, feature_size)
# (배치 수, 채널 수 * 이미지 너비 * 이미지 높이) ->
# (배치 수, 채널 수, 이미지 너비, 이미지 높이)
class Autoencoder(nn.Module): # encoder and decoder
def __init__(self):
super(Autoencoder, self).__init__()
k = 16
self.encoder = nn.Sequential(
nn.Conv2d(1, k, 3, stride=2),
nn.ReLU(),
nn.Conv2d(k, 2*k, 3, stride=2),
nn.ReLU(),
nn.Conv2d(2*k, 4*k, 3, stride=1),
nn.ReLU(),
Flatten(), # linear 들어가기 전에 flatten
nn.Linear(1024, 10),
nn.ReLU()
)
# ConvTranspose(2d)란?
# 입력 성분(Conv의 결과)을 출력 성분(Conv의 입력)으로 미분: 결과를 출력으로 미분
# 그 값을 입력 벡터와 곱해 출력 벡터를 산출한다.
# 출력 된 벡터는 행렬 형태로 변환한다.
# 디코더에서는 사이즈를 늘려줘야하기 때문에 convolution을 못함
# ConvTranspose2d 라는 upsampling 함수 사용 !!!!
self.decoder = nn.Sequential(
nn.Linear(10, 1024),
nn.ReLU(),
Deflatten(4*k), # linear 나오고 deflatten
nn.ConvTranspose2d(4*k, 2*k, 3, stride=1),
# (입력 채널 수, 출력 채널 수, 필터 크기, stride)
nn.ReLU(),
nn.ConvTranspose2d(2*k, k, 3, stride=2),
nn.ReLU(),
nn.ConvTranspose2d(k, 1, 3, stride=2,output_padding=1),
nn.Sigmoid()
)
def forward(self, x):
encoded = self.encoder(x)
decoded = self.decoder(encoded)
return decodedAutoencoder 모델의 구조를 살펴보면
model = Autoencoder().to(device)
model
# result
Autoencoder(
(encoder): Sequential(
(0): Conv2d(1, 16, kernel_size=(3, 3), stride=(2, 2))
(1): ReLU()
(2): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2))
(3): ReLU()
(4): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
(5): ReLU()
(6): Flatten() # linear 바로 전에 flatten
(7): Linear(in_features=1024, out_features=10, bias=True)
(8): ReLU()
)
(decoder): Sequential(
(0): Linear(in_features=10, out_features=1024, bias=True)
(1): ReLU()
(2): Deflatten() # linear 바로 이후 deflatten
(3): ConvTranspose2d(64, 32, kernel_size=(3, 3), stride=(1, 1))
(4): ReLU()
(5): ConvTranspose2d(32, 16, kernel_size=(3, 3), stride=(2, 2))
(6): ReLU()
(7): ConvTranspose2d(16, 1, kernel_size=(3, 3), stride=(2, 2), output_padding=(1, 1))
(8): Sigmoid()
)
)4. 시각화를 위한 함수 만들기
def normalize_output(img):
img = (img - img.min())/(img.max()-img.min())
return img
def check_plot():
with torch.no_grad():
for data in trainloader:
inputs = data[0].to(device)
outputs = model(inputs)
input_samples = inputs.permute(0,2,3,1).cpu().numpy() # 원래 이미지
reconstructed_samples = outputs.permute(0,2,3,1).cpu().numpy() # 생성 이미지
break # 배치 하나만 받고 for문 종료
#reconstructed_samples = normalize_output(reconstructed_samples) # 0~1사이로 변환
#input_samples = normalize_output(input_samples) # 0~1사이로 변환
columns = 10 # 시각화 전체 너비
rows = 5 # 시각화 전체 높이
fig=plt.figure(figsize=(columns, rows)) # figure 선언
# 원래 이미지 배치 크기 만큼 보여주기
for i in range(1, columns*rows+1):
img = input_samples[i-1]
fig.add_subplot(rows, columns, i)
plt.imshow(img.squeeze())
#plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plt.axis('off')
plt.show()
# 생성 이미지 배치 크기 만큼 보여주기
fig=plt.figure(figsize=(columns, rows))
for i in range(1, columns*rows+1):
img = reconstructed_samples[i-1]
fig.add_subplot(rows, columns, i)
plt.imshow(img.squeeze())
#plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plt.axis('off')
plt.show()5. 손실 함수 및 최적화 방법 정의
criterion = nn.MSELoss() # MSE 사용
optimizer = optim.Adam(model.parameters(), lr=1e-4)6. 모델 학습하기
for epoch in range(31):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs = data[0].to(device) # label이 없음
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(inputs, outputs) # 라벨 대신 입력 이미지와 출력 이미지를 비교
loss.backward()
optimizer.step()
running_loss += loss.item()
cost = running_loss / len(trainloader)
if epoch % 10 == 0:
print('[%d] loss: %.3f' %(epoch + 1, cost))
check_plot()7. Training Epoch별 시각화 결과
Epoch=1
Epoch=11
Epoch=21
Epoch=31
결론: epoch이 진행될수록 모델이 내뱉는 reconstructed image가 input 이미지와 유사해짐을 알 수 있음
*참고: Pytorch가 제공하는 CONVTRANSPOSE2D
convolution이 downsampling 기능이아면 convtranspose는 upsampling 기능을 제공하는 함수.
(아래와 같은 느낌으로)
pppanghyun