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

Future Prediction Training

Browse Source
This commit is contained in:
Si11ium
2019-09-29 11:50:38 +02:00
parent a70c9b7fef
commit 3e9ef013b3
8 changed files with 86 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
networks/modules.py
View File

@ -6,7 +6,7 @@ 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 torch.nn import Module, Linear, ReLU, Sigmoid, Dropout, GRU, Tanh
from torchvision.transforms import Normalize
from abc import ABC, abstractmethod
@ -33,8 +33,7 @@ class LightningModuleOverrides:
@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]),
return DataLoader(DataContainer(os.path.join('data', 'training'), self.size, self.step),
shuffle=True, batch_size=10000, num_workers=num_workers)
"""
@data_loader
@ -193,17 +192,23 @@ class Discriminator(Module):
class DecoderLinearStack(Module):
def __init__(self, out_shape):
def __init__(self, out_shape, use_norm=True):
super(DecoderLinearStack, self).__init__()
self.l1 = Linear(10, 100, bias=True)
self.norm1 = torch.nn.BatchNorm1d(100) if use_norm else False
self.l2 = Linear(100, out_shape, bias=True)
self.norm2 = torch.nn.BatchNorm1d(out_shape) if use_norm else False
self.activation = ReLU()
self.activation_out = Sigmoid()
self.activation_out = Tanh()
def forward(self, x):
tensor = self.l1(x)
if self.norm1:
tensor = self.norm1(tensor)
tensor = self.activation(tensor)
tensor = self.l2(tensor)
if self.norm2:
tensor = self.norm2(tensor)
tensor = self.activation_out(tensor)
return tensor
@ -213,62 +218,64 @@ class EncoderLinearStack(Module):
@property
def shape(self):
x = randn(self.features).unsqueeze(0)
x = torch.cat((x,x,x,x,x))
output = self(x)
return output.shape[1:]
def __init__(self, features=6, separated=False, use_bias=True):
def __init__(self, features=6, factor=10, use_bias=True, use_norm=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.l1 = Linear(self.features, self.features * factor, bias=use_bias)
self.l2 = Linear(self.features * factor, self.features * factor//2, bias=use_bias)
self.l3 = Linear(self.features * factor//2, factor, use_bias)
self.norm1 = torch.nn.BatchNorm1d(self.features * factor) if use_norm else False
self.norm2 = torch.nn.BatchNorm1d(self.features * factor//2) if use_norm else False
self.norm3 = torch.nn.BatchNorm1d(factor) if use_norm else False
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.l1(x)
if self.norm1:
tensor = self.norm1(tensor)
tensor = self.activation(tensor)
tensor = self.l2(tensor)
if self.norm2:
tensor = self.norm2(tensor)
tensor = self.activation(tensor)
tensor = self.l3(tensor)
if self.norm3:
tensor = self.norm3(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):
def __init__(self, lat_dim, variational=False, use_dense=True, features=6, use_norm=True):
self.lat_dim = lat_dim
self.features = features
self.lstm_cells = 10
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.l_stack = TimeDistributed(EncoderLinearStack(features, use_norm=use_norm)) if use_dense else False
self.gru = GRU(10 if use_dense else self.features, self.lstm_cells, batch_first=True)
self.filter = RNNOutputFilter(only_last=True)
self.norm = torch.nn.BatchNorm1d(self.lstm_cells) if use_norm else False
if variational:
self.mu = Linear(10, self.lat_dim)
self.logvar = Linear(10, self.lat_dim)
self.mu = Linear(self.lstm_cells, self.lat_dim)
self.logvar = Linear(self.lstm_cells, self.lat_dim)
else:
self.lat_dim_layer = Linear(10, self.lat_dim)
self.lat_dim_layer = Linear(self.lstm_cells, 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.norm:
tensor = self.norm(tensor)
if self.variational:
tensor = self.mu(tensor), self.logvar(tensor)
else:
@ -316,17 +323,20 @@ class PoolingEncoder(Module):
class Decoder(Module):
def __init__(self, latent_dim, *args, variational=False):
def __init__(self, latent_dim, *args, lstm_cells=10, use_norm=True, variational=False):
self.variational = variational
super(Decoder, self).__init__()
self.g = GRU(latent_dim, 10, batch_first=True)
self.gru = GRU(latent_dim, lstm_cells, batch_first=True)
self.norm = TimeDistributed(torch.nn.BatchNorm1d(lstm_cells) if use_norm else False)
self.filter = RNNOutputFilter()
self.l_stack = TimeDistributed(DecoderLinearStack(*args))
self.l_stack = TimeDistributed(DecoderLinearStack(*args, use_norm=use_norm))
pass
def forward(self, x):
tensor = self.g(x)
tensor = self.gru(x)
tensor = self.filter(tensor)
if self.norm:
tensor = self.norm(tensor)
tensor = self.l_stack(tensor)
return tensor