steffen/ae_toolbox_torch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

transition

Browse Source
This commit is contained in:
Si11ium
2021-02-01 09:59:56 +01:00
parent 4c489237d7
commit 578727d043
35 changed files with 177 additions and 305 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
networks/modules.py
View File

@ -5,9 +5,7 @@ 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, Tanh
from torchvision.transforms import Normalize
from abc import ABC, abstractmethod
@ -27,21 +25,12 @@ class LightningModuleOverrides:
def name(self):
return self.__class__.__name__
def forward(self, x):
return self.network.forward(x)
@data_loader
@pl.data_loader
def train_dataloader(self):
num_workers = 0 # os.cpu_count() // 2
num_workers = 0 # os.cpu_count() // 2
return DataLoader(DataContainer(os.path.join('data', 'training'), self.size, self.step),
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):
@ -56,53 +45,6 @@ class AbstractNeuralNetwork(Module):
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
@abstractmethod
def configure_optimizers(self):
# REQUIRED
raise NotImplementedError
@pl.data_loader
def train_dataloader(self):
# REQUIRED
raise NotImplementedError
"""
def validation_step(self, batch, batch_nb):
# OPTIONAL
pass
def validation_end(self, outputs):
# OPTIONAL
pass
@pl.data_loader
def val_dataloader(self):
# OPTIONAL
pass
@pl.data_loader
def test_dataloader(self):
# OPTIONAL
pass
"""
#######################
# Utility Modules
class TimeDistributed(Module):
@ -167,12 +109,14 @@ class AvgDimPool(Module):
# Generators, Decoders, Encoders, Discriminators
class Discriminator(Module):
def __init__(self, latent_dim, features, dropout=.0, activation=ReLU):
def __init__(self, latent_dim, features, dropout=.0, activation=ReLU, use_norm=False):
super(Discriminator, self).__init__()
self.features = features
self.latent_dim = latent_dim
self.l1 = Linear(self.latent_dim, self.features * 10)
self.norm1 = torch.nn.BatchNorm1d(self.features * 10) if use_norm else False
self.l2 = Linear(self.features * 10, self.features * 20)
self.norm2 = torch.nn.BatchNorm1d(self.features * 20) if use_norm else False
self.lout = Linear(self.features * 20, 1)
self.dropout = Dropout(dropout)
self.activation = activation()
@ -180,9 +124,15 @@ class Discriminator(Module):
def forward(self, x, **kwargs):
tensor = self.l1(x)
tensor = self.dropout(self.activation(tensor))
tensor = self.dropout(tensor)
if self.norm1:
tensor = self.norm1(tensor)
tensor = self.activation(tensor)
tensor = self.l2(tensor)
tensor = self.dropout(self.activation(tensor))
tensor = self.dropout(tensor)
if self.norm2:
tensor = self.norm2(tensor)
tensor = self.activation(tensor)
tensor = self.lout(tensor)
tensor = self.sigmoid(tensor)
return tensor
@ -296,13 +246,13 @@ class AttentionEncoder(Module):
class PoolingEncoder(Module):
def __init__(self, lat_dim, variational=False):
def __init__(self, lat_dim, variational=False, use_norm=True):
self.lat_dim = lat_dim
self.variational = variational
super(PoolingEncoder, self).__init__()
self.p = AvgDimPool()
self.l = EncoderLinearStack()
self.l = EncoderLinearStack(use_norm=use_norm)
if variational:
self.mu = Linear(self.l.shape, self.lat_dim)
self.logvar = Linear(self.l.shape, self.lat_dim)