#!/usr/bin/env python3
# coding: utf-8
# source:
# https://www.kaggle.com/code/weka511/autoencoder-implementation-in-pytorch
from matplotlib.pyplot import close, figure, imshow, savefig, show, title
from matplotlib.lines import Line2D
from os.path import join # pathname manipulation
from random import sample # random sampling of lists
from re import split # regular expression (strings)
from torch import device, no_grad
from torch.cuda import is_available
from torch.nn import Linear, Module, MSELoss, ReLU, Sequential, Sigmoid
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor
from torchvision.utils import make_grid
#
# Hyperparameters
#
# 1. The sizes of the encoder layers are taken from
# [Reducing the Dimensionality of Data with Neural Networks
# G. E. Hinton and R. R. Salakhutdinov]
# (https://www.cs.toronto.edu/~hinton/science.pdf)
# 2. The learning rate was optimized by trial and error.
# The error rates are plotted here
# (https://github.com/weka511/learn/issues/26)
ENCODER = [28*28,400,200,100,50,25,6] # sizes of encoder layers
DECODER = [] # Decoder layers will be a mirror image of encoder
LR = 0.001 # Learning rate
N = 32 # Number of epochs
#
# The Autoencoder class
#
# The latest version of this class can be found at
# [github](https://github.com/weka511/learn/blob/master/ae.py)
class AutoEncoder(Module):
'''A class that implements an AutoEncoder
'''
@staticmethod
def get_non_linearity(params):
'''Determine which non linearity is to be used for both
encoder and decoder'''
def get_one(param):
'''Determine which non linearity is to be used for
either encoder or decoder'''
param = param.lower()
if param=='relu': return ReLU()
if param=='sigmoid': return Sigmoid()
return None
decoder_non_linearity = get_one(params[0])
encoder_non_linearity = \
getnl(params[a]) if len(params)>1 else decoder_non_linearity
return encoder_non_linearity, decoder_non_linearity
@staticmethod
def build_layer(sizes,
non_linearity = None):
'''Construct encoder or decoder as a Sequential of Linear
labels, with or without non-linearities
Positional arguments:
sizes List of sizes for each Linear Layer
Keyword arguments:
non_linearity Object used to introduce non-linearity between layers
'''
linears = [Linear(m,n) for m,n in zip(sizes[:-1],sizes[1:])]
if non_linearity==None:
return Sequential(*linears)
else:
return Sequential(*[item for pair in [(layer,non_linearity) \
for layer in linears] for item in pair])
def __init__(self,
encoder_sizes = [28*28,400,200,100,50,25,6],
encoder_non_linearity = ReLU(inplace=True),
decoder_sizes = [],
decoder_non_linearity = ReLU(inplace=True)):
'''Keyword arguments:
encoder_sizes List of sizes for each Linear Layer in encoder
encoder_non_linearity Non-linearity between encoder layers
decoder_sizes List of sizes for each Linear Layer in decoder
decoder_non_linearity Non-linearity between decoder layers
'''
super().__init__()
self.encoder_sizes = encoder_sizes
self.decoder_sizes = encoder_sizes[::-1] if len(decoder_sizes)==0 \
else decoder_sizes
self.encoder = AutoEncoder.build_layer(self.encoder_sizes,
non_linearity = encoder_non_linearity)
self.decoder = AutoEncoder.build_layer(self.decoder_sizes,
non_linearity = decoder_non_linearity)
self.encodeBool = True
self.decodeBool = True
def forward(self, x):
'''Propagate value through network
Computation is controlled by self.encodeBool and self.decodeBool
'''
if self.encodeBool:
x = self.encoder(x)
if self.decodeBool:
x = self.decoder(x)
return x
def n_encoded(self):
return self.encoder_sizes[-1]
#
# Function to train network
#
def train(loader, model, optimizer, criterion, N = 25, dev = 'cpu'):
'''Train network
Parameters:
loader Used to get data
model Model to be trained
optimizer Used to minimze errors
criterion Used to compute errors
Keyword parameters:
N Number of epochs
dev Device - cpu or cuda
'''
Losses = []
for epoch in range(N):
loss = 0
for batch_features, _ in loader:
batch_features = batch_features.view(-1, 784).to(dev)
optimizer.zero_grad()
outputs = model(batch_features)
train_loss = criterion(outputs, batch_features)
train_loss.backward()
optimizer.step()
loss += train_loss.item()
Losses.append(loss / len(loader))
print(f'epoch : {epoch+1}/{N}, loss = {Losses[-1]:.6f}')
return Losses
#
# Initialize network and data, and prepare to train
#
# This is proably a suboptimal way to load the MNIST dataset,
# but it will do for this example.
#
dev = device("cuda" if is_available() else "cpu")
encoder_non_linearity,decoder_non_linearity = AutoEncoder.get_non_linearity(['relu'])
model = AutoEncoder(encoder_sizes = ENCODER,
encoder_non_linearity = encoder_non_linearity,
decoder_non_linearity = decoder_non_linearity,
decoder_sizes = DECODER).to(dev)
optimizer = Adam(model.parameters(),
lr = LR)
criterion = MSELoss()
transform = Compose([ToTensor()])
train_dataset = MNIST(root="~/torch_datasets",
train = True,
transform = transform,
download = True)
test_dataset = MNIST(root="~/torch_datasets",
train = False,
transform = transform,
download = True)
train_loader = DataLoader(train_dataset,
batch_size = 128,
shuffle = True,
num_workers = 4)
test_loader = DataLoader(test_dataset,
batch_size = 32,
shuffle = False,
num_workers = 4)
#
# Train network
#
Losses = train(train_loader,model,optimizer,criterion, N = N, dev = dev)
def reconstruct(loader,model,criterion,
N = 25,
prefix = 'test',
show = False,
figs = './figs',
n_images = -1):
'''Reconstruct images from encoding
Parameters:
loader
model
Keyword Parameters:
N Number of epochs used for training (used in image title only)
prefix Prefix file names with this string
show Used to display images
figs Directory for storing images
'''
def plot(original=None,decoded=None):
'''Plot original images and decoded images'''
fig = figure(figsize=(10,10))
ax = fig.subplots(nrows=2)
ax[0].imshow(make_grid(original.view(-1,1,28,28)).permute(1, 2, 0))
ax[0].set_title('Raw images')
scaled_decoded = decoded/decoded.max()
ax[1].imshow(make_grid(scaled_decoded.view(-1,1,28,28)).permute(1, 2, 0))
ax[1].set_title(f'Reconstructed images after {N} epochs')
savefig(join(figs,f'{prefix}-comparison-{i}'))
if not show:
close (fig)
samples = [] if n_images==-1 else sample(range(len(loader)//loader.batch_size),
k = n_images)
loss = 0.0
with no_grad():
for i,(batch_features, _) in enumerate(loader):
batch_features = batch_features.view(-1, 784).to(dev)
outputs = model(batch_features)
test_loss = criterion(outputs, batch_features)
loss += test_loss.item()
if len(samples)==0 or i in samples:
plot(original=batch_features,
decoded=outputs)
return loss
#
# Compare output layer with Inputs,
# to get an idea of the quality of the encoding
#
test_loss = reconstruct(test_loader,model,criterion,
N = N,
show = True,
figs = '.',
n_images = 5,
prefix = 'foo')
def plot_losses(Losses,
lr = 0.001,
encoder = [],
decoder = [],
encoder_nonlinearity = None,
decoder_nonlinearity = None,
N = 25,
show = False,
figs = './figs',
prefix = 'ae',
test_loss = 0):
'''Plot curve of training losses'''
fig = figure(figsize=(10,10))
ax = fig.subplots()
ax.plot(Losses)
ax.set_ylim(bottom=0)
ax.set_title(f'Training Losses after {N} epochs')
ax.set_ylabel('MSELoss')
ax.text(0.95, 0.95, '\n'.join([f'lr = {lr}',
f'encoder = {encoder}',
f'decoder = {decoder}',
f'encoder nonlinearity = {encoder_nonlinearity}',
f'decoder nonlinearity = {decoder_nonlinearity}',
f'test loss = {test_loss:.3f}'
]),
transform = ax.transAxes,
fontsize = 14,
verticalalignment = 'top',
horizontalalignment = 'right',
bbox = dict(boxstyle = 'round',
facecolor = 'wheat',
alpha = 0.5))
savefig(join(figs,f'{prefix}-losses'))
if not show:
close (fig)
plot_losses(Losses,
lr = LR,
encoder = model.encoder_sizes,
decoder = model.decoder_sizes,
encoder_nonlinearity = encoder_non_linearity,
decoder_nonlinearity = decoder_non_linearity,
N = N,
show = True,
figs = '.',
prefix = 'foo',
test_loss = test_loss)
def plot_encoding(loader,model,
figs = './figs',
dev = 'cpu',
colours = [],
show = False,
prefix = 'ae'):
'''Plot the encoding layer
Since this is multi,dimensional, we will break it into 2D plots
'''
def extract_batch(batch_features, labels,index):
'''Extract xs, ys, and colours for one batch'''
batch_features = batch_features.view(-1, 784).to(dev)
encoded = model(batch_features).tolist()
return list(zip(*([encoded[k][2*index] for k in range(len(labels))],
[encoded[k][2*index+1] for k in range(len(labels))],
[colours[labels.tolist()[k]] for k in range(len(labels))])))
save_decode = model.decodeBool
model.decodeBool = False
with no_grad():
fig = figure(figsize=(10,10))
ax = fig.subplots(nrows=2,ncols=2)
for i in range(2):
for j in range(2):
if i==1 and j==1: break
index = 2*i + j
if 2*index+1 < model.n_encoded():
xs,ys,cs = tuple(zip(*[xyc for batch_features, labels in loader for xyc in extract_batch(batch_features, labels,index)]))
ax[i][j].set_title(f'{2*index}-{2*index+1}')
ax[i][j].scatter(xs,ys,c=cs,s=1)
ax[0][0].legend(handles=[Line2D([], [],
color = colours[k],
marker = 's',
ls = '',
label = f'{k}') for k in range(10)])
savefig(join(figs,f'{prefix}-encoding'))
if not show:
close (fig)
model.decode = save_decode
#
# Plot encoded data
#
# The encoding shows that the images for most digits are separated.
# It also suggest that the encoded data clouls have been made to
# live in a 5 dimensional manifold instead of needind 6.
#
plot_encoding(test_loader,model,
show = True,
colours = ['xkcd:purple',
'xkcd:green',
'xkcd:blue',
'xkcd:pink',
'xkcd:brown',
'xkcd:red',
'xkcd:magenta',
'xkcd:yellow',
'xkcd:light teal',
'xkcd:puke'],
figs = '.',
prefix = 'foo')
$\qquad$
