336 lines
9.7 KiB
Python
336 lines
9.7 KiB
Python
import os
|
|
from operator import mul
|
|
from functools import reduce
|
|
|
|
import torch
|
|
from torch import randn
|
|
import pytorch_lightning as pl
|
|
from pytorch_lightning import data_loader
|
|
from torch.nn import Module, Linear, ReLU, Sigmoid, Dropout, GRU
|
|
from torchvision.transforms import Normalize
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
#######################
|
|
# Abstract NN Class & Lightning Module
|
|
from torch.utils.data import DataLoader
|
|
|
|
from dataset import DataContainer
|
|
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
|
|
class LightningModuleOverrides:
|
|
|
|
@property
|
|
def name(self):
|
|
return self.__class__.__name__
|
|
|
|
def forward(self, x):
|
|
return self.network.forward(x)
|
|
|
|
@data_loader
|
|
def tng_dataloader(self):
|
|
num_workers = 0 # os.cpu_count() // 2
|
|
return DataLoader(DataContainer(os.path.join('data', 'training'),
|
|
self.size, self.step, transforms=[Normalize]),
|
|
shuffle=True, batch_size=10000, num_workers=num_workers)
|
|
"""
|
|
@data_loader
|
|
def val_dataloader(self):
|
|
num_workers = 0 # os.cpu_count() // 2
|
|
return DataLoader(DataContainer(os.path.join('data', 'validation'), self.size, self.step),
|
|
shuffle=True, batch_size=100, num_workers=num_workers)
|
|
"""
|
|
|
|
|
|
class AbstractNeuralNetwork(Module):
|
|
|
|
@property
|
|
def name(self):
|
|
return self.__class__.__name__
|
|
|
|
def __init__(self):
|
|
super(AbstractNeuralNetwork, self).__init__()
|
|
|
|
def forward(self, batch):
|
|
pass
|
|
|
|
|
|
######################
|
|
# Abstract Network class following the Lightning Syntax
|
|
class LightningModule(pl.LightningModule, ABC):
|
|
|
|
def __init__(self):
|
|
super(LightningModule, self).__init__()
|
|
|
|
@abstractmethod
|
|
def forward(self, x):
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def training_step(self, batch, batch_nb):
|
|
# REQUIRED
|
|
raise NotImplementedError
|
|
|
|
def validation_step(self, batch, batch_nb):
|
|
# OPTIONAL
|
|
pass
|
|
|
|
def validation_end(self, outputs):
|
|
# OPTIONAL
|
|
pass
|
|
|
|
@abstractmethod
|
|
def configure_optimizers(self):
|
|
# REQUIRED
|
|
raise NotImplementedError
|
|
|
|
@pl.data_loader
|
|
def tng_dataloader(self):
|
|
# REQUIRED
|
|
raise NotImplementedError
|
|
# return DataLoader(MNIST(os.getcwd(), train=True, download=True,
|
|
# transform=transforms.ToTensor()), batch_size=32)
|
|
|
|
"""
|
|
@pl.data_loader
|
|
def val_dataloader(self):
|
|
# OPTIONAL
|
|
pass
|
|
|
|
@pl.data_loader
|
|
def test_dataloader(self):
|
|
# OPTIONAL
|
|
pass
|
|
"""
|
|
|
|
#######################
|
|
# Utility Modules
|
|
class TimeDistributed(Module):
|
|
def __init__(self, module, batch_first=True):
|
|
super(TimeDistributed, self).__init__()
|
|
self.module = module
|
|
self.batch_first = batch_first
|
|
|
|
def forward(self, x):
|
|
|
|
if len(x.size()) <= 2:
|
|
return self.module(x)
|
|
|
|
# Squash samples and timesteps into a single axis
|
|
x_reshape = x.contiguous().view(-1, x.size(-1)) # (samples * timesteps, input_size)
|
|
|
|
y = self.module(x_reshape)
|
|
|
|
# We have to reshape Y
|
|
if self.batch_first:
|
|
y = y.contiguous().view(x.size(0), -1, y.size(-1)) # (samples, timesteps, output_size)
|
|
else:
|
|
y = y.view(-1, x.size(1), y.size(-1)) # (timesteps, samples, output_size)
|
|
|
|
return y
|
|
|
|
|
|
class Repeater(Module):
|
|
def __init__(self, shape):
|
|
super(Repeater, self).__init__()
|
|
self.shape = shape
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
x = x.unsqueeze(-2)
|
|
return x.expand(self.shape)
|
|
|
|
|
|
class RNNOutputFilter(Module):
|
|
|
|
def __init__(self, return_output=True, only_last=False):
|
|
super(RNNOutputFilter, self).__init__()
|
|
self.only_last = only_last
|
|
self.return_output = return_output
|
|
|
|
def forward(self, x: tuple):
|
|
outputs, hidden = x
|
|
out = outputs if self.return_output else hidden
|
|
return out if not self.only_last else out[:, -1, :]
|
|
|
|
|
|
class AvgDimPool(Module):
|
|
|
|
def __init__(self):
|
|
super(AvgDimPool, self).__init__()
|
|
|
|
def forward(self, x):
|
|
return x.mean(-2)
|
|
|
|
|
|
#######################
|
|
# Network Modules
|
|
# Generators, Decoders, Encoders, Discriminators
|
|
class Discriminator(Module):
|
|
|
|
def __init__(self, latent_dim, features, dropout=.0, activation=ReLU):
|
|
super(Discriminator, self).__init__()
|
|
self.features = features
|
|
self.latent_dim = latent_dim
|
|
self.l1 = Linear(self.latent_dim, self.features * 10)
|
|
self.l2 = Linear(self.features * 10, self.features * 20)
|
|
self.lout = Linear(self.features * 20, 1)
|
|
self.dropout = Dropout(dropout)
|
|
self.activation = activation()
|
|
self.sigmoid = Sigmoid()
|
|
|
|
def forward(self, x, **kwargs):
|
|
tensor = self.l1(x)
|
|
tensor = self.dropout(self.activation(tensor))
|
|
tensor = self.l2(tensor)
|
|
tensor = self.dropout(self.activation(tensor))
|
|
tensor = self.lout(tensor)
|
|
tensor = self.sigmoid(tensor)
|
|
return tensor
|
|
|
|
|
|
class DecoderLinearStack(Module):
|
|
|
|
def __init__(self, out_shape):
|
|
super(DecoderLinearStack, self).__init__()
|
|
self.l1 = Linear(10, 100, bias=True)
|
|
self.l2 = Linear(100, out_shape, bias=True)
|
|
self.activation = ReLU()
|
|
self.activation_out = Sigmoid()
|
|
|
|
def forward(self, x):
|
|
tensor = self.l1(x)
|
|
tensor = self.activation(tensor)
|
|
tensor = self.l2(tensor)
|
|
tensor = self.activation_out(tensor)
|
|
return tensor
|
|
|
|
|
|
class EncoderLinearStack(Module):
|
|
|
|
@property
|
|
def shape(self):
|
|
x = randn(self.features).unsqueeze(0)
|
|
output = self(x)
|
|
return output.shape[1:]
|
|
|
|
def __init__(self, features=6, separated=False, use_bias=True):
|
|
super(EncoderLinearStack, self).__init__()
|
|
# FixMe: Get Hardcoded shit out of here
|
|
self.separated = separated
|
|
self.features = features
|
|
if self.separated:
|
|
self.l1s = [Linear(1, 10, bias=use_bias) for _ in range(self.features)]
|
|
self.l2s = [Linear(10, 5, bias=use_bias) for _ in range(self.features)]
|
|
else:
|
|
self.l1 = Linear(self.features, self.features * 10, bias=use_bias)
|
|
self.l2 = Linear(self.features * 10, self.features * 5, bias=use_bias)
|
|
self.l3 = Linear(self.features * 5, 10, use_bias)
|
|
self.activation = ReLU()
|
|
|
|
def forward(self, x):
|
|
if self.separated:
|
|
x = x.unsqueeze(-1)
|
|
tensors = [self.l1s[idx](x[:, idx, :]) for idx in range(len(self.l1s))]
|
|
tensors = [self.activation(tensor) for tensor in tensors]
|
|
tensors = [self.l2s[idx](tensors[idx]) for idx in range(len(self.l2s))]
|
|
tensors = [self.activation(tensor) for tensor in tensors]
|
|
tensor = torch.cat(tensors, dim=-1)
|
|
else:
|
|
tensor = self.l1(x)
|
|
tensor = self.activation(tensor)
|
|
tensor = self.l2(tensor)
|
|
tensor = self.l3(tensor)
|
|
tensor = self.activation(tensor)
|
|
return tensor
|
|
|
|
|
|
class Encoder(Module):
|
|
|
|
def __init__(self, lat_dim, variational=False, separate_features=False, with_dense=True, features=6):
|
|
self.lat_dim = lat_dim
|
|
self.features = features
|
|
self.variational = variational
|
|
super(Encoder, self).__init__()
|
|
self.l_stack = TimeDistributed(EncoderLinearStack(separated=separate_features,
|
|
features=features)) if with_dense else False
|
|
self.gru = GRU(10 if with_dense else self.features, 10, batch_first=True)
|
|
self.filter = RNNOutputFilter(only_last=True)
|
|
if variational:
|
|
self.mu = Linear(10, self.lat_dim)
|
|
self.logvar = Linear(10, self.lat_dim)
|
|
else:
|
|
self.lat_dim_layer = Linear(10, self.lat_dim)
|
|
|
|
def forward(self, x):
|
|
if self.l_stack:
|
|
x = self.l_stack(x)
|
|
tensor = self.gru(x)
|
|
tensor = self.filter(tensor)
|
|
if self.variational:
|
|
tensor = self.mu(tensor), self.logvar(tensor)
|
|
else:
|
|
tensor = self.lat_dim_layer(tensor)
|
|
return tensor
|
|
|
|
|
|
class AttentionEncoder(Module):
|
|
|
|
def __init__(self):
|
|
super(AttentionEncoder, self).__init__()
|
|
self.l_stack = TimeDistributed(EncoderLinearStack())
|
|
|
|
def forward(self, x):
|
|
tensor = self.l_stack(x)
|
|
torch.bmm() # TODO Add Attention here
|
|
|
|
return tensor
|
|
|
|
|
|
class PoolingEncoder(Module):
|
|
|
|
def __init__(self, lat_dim, variational=False):
|
|
self.lat_dim = lat_dim
|
|
self.variational = variational
|
|
|
|
super(PoolingEncoder, self).__init__()
|
|
self.p = AvgDimPool()
|
|
self.l = EncoderLinearStack()
|
|
if variational:
|
|
self.mu = Linear(self.l.shape, self.lat_dim)
|
|
self.logvar = Linear(self.l.shape, self.lat_dim)
|
|
else:
|
|
self.lat_dim_layer = Linear(reduce(mul, self.l.shape), self.lat_dim)
|
|
|
|
def forward(self, x):
|
|
tensor = self.p(x)
|
|
tensor = self.l(tensor)
|
|
if self.variational:
|
|
tensor = self.mu(tensor), self.logvar(tensor)
|
|
else:
|
|
tensor = self.lat_dim_layer(tensor)
|
|
return tensor
|
|
|
|
|
|
class Decoder(Module):
|
|
|
|
def __init__(self, latent_dim, *args, variational=False):
|
|
self.variational = variational
|
|
super(Decoder, self).__init__()
|
|
self.g = GRU(latent_dim, 10, batch_first=True)
|
|
self.filter = RNNOutputFilter()
|
|
self.l_stack = TimeDistributed(DecoderLinearStack(*args))
|
|
pass
|
|
|
|
def forward(self, x):
|
|
tensor = self.g(x)
|
|
tensor = self.filter(tensor)
|
|
tensor = self.l_stack(tensor)
|
|
return tensor
|
|
|
|
|
|
if __name__ == '__main__':
|
|
raise PermissionError('Get out of here - never run this module')
|