Network.py

# -*- coding: utf-8 -*-
"""
Created on Fri Nov 18 12:17:53 2022

@author: srpv

- Network architecture
- LSTM
- Hidden dimension
- No of recurrent layers
- Bottle neck dimension


"""

from torch import nn, optim
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #Setting Up CUDA

class Encoder(nn.Module):

  def __init__(self, seq_len, n_features, embedding_dim=128):
    super(Encoder, self).__init__()

    self.seq_len, self.n_features = seq_len, n_features
    self.embedding_dim, self.hidden_dim = embedding_dim, 2 * embedding_dim

    self.rnn1 = nn.LSTM(
      input_size=n_features,
      hidden_size=self.hidden_dim,
      num_layers=1,
      batch_first=True
    )
    
    self.rnn2 = nn.LSTM(
      input_size=self.hidden_dim,
      hidden_size=embedding_dim,
      num_layers=1,
      batch_first=True
    )

  def forward(self, x):
      
    x = x.reshape((1, self.seq_len, self.n_features))
    x, (_, _) = self.rnn1(x) 
    x, (hidden_n, _) = self.rnn2(x)
    x=hidden_n.reshape((self.n_features, self.embedding_dim))
     
    return x



#%%

class Decoder(nn.Module):

  def __init__(self, seq_len, input_dim=128, n_features=1):
    super(Decoder, self).__init__()

    self.seq_len, self.input_dim = seq_len, input_dim
    self.hidden_dim, self.n_features = 2 * input_dim, n_features

    self.rnn1 = nn.LSTM(
      input_size=input_dim,
      hidden_size=input_dim,
      num_layers=1,
      batch_first=True
    )

    self.rnn2 = nn.LSTM(
      input_size=input_dim,
      hidden_size=self.hidden_dim,
      num_layers=1,
      batch_first=True
    )

    self.output_layer = nn.Linear(self.hidden_dim, n_features)

  def forward(self, x):
      
    
    x = x.repeat(self.seq_len, self.n_features)
    x = x.reshape((self.n_features, self.seq_len, self.input_dim))
    x, (hidden_n, cell_n) = self.rnn1(x)
    x, (hidden_n, cell_n) = self.rnn2(x)
    x = x.reshape((self.seq_len, self.hidden_dim))
    x=self.output_layer(x)


    return x