!pip install torch torchvision
Requirement already satisfied: torch in /usr/local/lib/python3.6/dist-packages (1.1.0) Requirement already satisfied: torchvision in /usr/local/lib/python3.6/dist-packages (0.3.0) Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch) (1.16.4) Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from torchvision) (1.12.0) Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision) (4.3.0) Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from pillow>=4.1.1->torchvision) (0.46)
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
batch_size = 256
learning_rate = 0.0002
num_epoch = 5
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)
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)
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
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
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)
cuda:0
# ¸ðµ¨À» ºÒ·¯¿À´Â ¹æ¹ýÀÔ´Ï´Ù.
# Å©°Ô µÎ°¡Áö ¹æ¹ýÀÌ Àִµ¥ ¿©±â »ç¿ëµÈ ¹æ¹ýÀº Á» ´Ü¼øÇÑ ¹æ¹ýÀÔ´Ï´Ù.
# 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)
--------model not restored--------
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()
torch.Size([256, 28, 28])
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)
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()
torch.Size([256, 28, 28])