Convolutional Autoencoder¶
- MNIST dataset
- Convolutional Neural Network
- 2 hidden layers
In [1]:
!pip install torch torchvision
In [0]:
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.init as init
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
%matplotlib inline
2) Set hyperparameters¶
In [0]:
batch_size = 256
learning_rate = 0.0002
num_epoch = 5
In [0]:
mnist_train = dset.MNIST("./", train=True, transform=transforms.ToTensor(), target_transform=None, download=True)
mnist_test = dset.MNIST("./", train=False, transform=transforms.ToTensor(), target_transform=None, download=True)
2) Set DataLoader¶
In [0]:
train_loader = torch.utils.data.DataLoader(mnist_train,batch_size=batch_size, shuffle=True,num_workers=2,drop_last=True)
test_loader = torch.utils.data.DataLoader(mnist_test,batch_size=batch_size, shuffle=False,num_workers=2,drop_last=True)
In [0]:
class Encoder(nn.Module):
def __init__(self):
super(Encoder,self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1,16,3,padding=1), # batch x 16 x 28 x 28
nn.ReLU(),
nn.BatchNorm2d(16),
nn.Conv2d(16,32,3,padding=1), # batch x 32 x 28 x 28
nn.ReLU(),
nn.BatchNorm2d(32),
nn.Conv2d(32,64,3,padding=1), # batch x 32 x 28 x 28
nn.ReLU(),
nn.BatchNorm2d(64),
nn.MaxPool2d(2,2) # batch x 64 x 14 x 14
)
self.layer2 = nn.Sequential(
nn.Conv2d(64,128,3,padding=1), # batch x 64 x 14 x 14
nn.ReLU(),
nn.BatchNorm2d(128),
nn.MaxPool2d(2,2),
nn.Conv2d(128,256,3,padding=1), # batch x 64 x 7 x 7
nn.ReLU()
)
def forward(self,x):
out = self.layer1(x)
out = self.layer2(out)
out = out.view(batch_size, -1)
return out
In [0]:
class Decoder(nn.Module):
def __init__(self):
super(Decoder,self).__init__()
self.layer1 = nn.Sequential(
nn.ConvTranspose2d(256,128,3,2,1,1), # batch x 128 x 14 x 14
nn.ReLU(),
nn.BatchNorm2d(128),
nn.ConvTranspose2d(128,64,3,1,1), # batch x 64 x 14 x 14
nn.ReLU(),
nn.BatchNorm2d(64)
)
self.layer2 = nn.Sequential(
nn.ConvTranspose2d(64,16,3,1,1), # batch x 16 x 14 x 14
nn.ReLU(),
nn.BatchNorm2d(16),
nn.ConvTranspose2d(16,1,3,2,1,1), # batch x 1 x 28 x 28
nn.ReLU()
)
def forward(self,x):
out = x.view(batch_size,256,7,7)
out = self.layer1(out)
out = self.layer2(out)
return out
2) Loss func & Optimizer¶
In [8]:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
encoder = Encoder().to(device)
decoder = Decoder().to(device)
# ÀÎÄÚ´õ µðÄÚ´õÀÇ ÆĶó¹ÌÅ͸¦ µ¿½Ã¿¡ ÇнÀ½ÃÅ°±â À§ÇØ À̸¦ ¹´Â ¹æ¹ýÀÔ´Ï´Ù.
parameters = list(encoder.parameters())+ list(decoder.parameters())
loss_func = nn.MSELoss()
optimizer = torch.optim.Adam(parameters, lr=learning_rate)
5. Train¶
In [9]:
# ¸ðµ¨À» ºÒ·¯¿À´Â ¹æ¹ýÀÔ´Ï´Ù.
# Å©°Ô µÎ°¡Áö ¹æ¹ýÀÌ Àִµ¥ ¿©±â »ç¿ëµÈ ¹æ¹ýÀº Á» ´Ü¼øÇÑ ¹æ¹ýÀÔ´Ï´Ù.
# https://pytorch.org/tutorials/beginner/saving_loading_models.html
try:
encoder, decoder = torch.load('./model/conv_autoencoder.pkl')
print("\n--------model restored--------\n")
except:
print("\n--------model not restored--------\n")
pass
for i in range(num_epoch):
for j,[image,label] in enumerate(train_loader):
optimizer.zero_grad()
image = image.to(device)
output = encoder(image)
output = decoder(output)
loss = loss_func(output,image)
loss.backward()
optimizer.step()
if j % 10 == 0:
# ¸ðµ¨ ÀúÀåÇÏ´Â ¹æ¹ý
# ÀÌ ¿ª½Ã Å©°Ô µÎ°¡Áö ¹æ¹ýÀÌ Àִµ¥ ¿©±â »ç¿ëµÈ ¹æ¹ýÀº Á» ´Ü¼øÇÑ ¹æ¹ýÀÔ´Ï´Ù.
# https://pytorch.org/tutorials/beginner/saving_loading_models.html
torch.save([encoder,decoder],'./model/conv_autoencoder.pkl')
print(loss)
6. Check with Train Image¶
In [10]:
out_img = torch.squeeze(output.cpu().data)
print(out_img.size())
for i in range(5):
plt.subplot(1,2,1)
plt.imshow(torch.squeeze(image[i]).cpu().numpy(),cmap='gray')
plt.subplot(1,2,2)
plt.imshow(out_img[i].numpy(),cmap='gray')
plt.show()
In [0]:
with torch.no_grad():
for j,[image,label] in enumerate(test_loader):
image = image.to(device)
output = encoder(image)
output = decoder(output)
if j % 10 == 0:
print(loss)
In [12]:
out_img = torch.squeeze(output.cpu().data)
print(out_img.size())
for i in range(5):
plt.subplot(1,2,1)
plt.imshow(torch.squeeze(image[i]).cpu().numpy(),cmap='gray')
plt.subplot(1,2,2)
plt.imshow(out_img[i].numpy(),cmap='gray')
plt.show()
In [0]: