CryptoModel.py

import torch.nn as nn
import torch

class CryptoModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, n_lyrs = 1, do = .05, device = "cpu"):
        """Initialize the network architecture
        Args:
            input_dim (int): [Number of time lags to look at for current prediction]
            hidden_dim (int): [The dimension of RNN output]
            n_lyrs (int, optional): [Number of stacked RNN layers]. Defaults to 1.
            do (float, optional): [Dropout for regularization]. Defaults to .05.
        """
        super(CryptoModel, self).__init__()

        self.ip_dim = input_dim
        self.hidden_dim = hidden_dim
        self.n_layers = n_lyrs
        self.dropout = do
        self.device = device

        self.rnn = nn.LSTM(input_size = input_dim, hidden_size = hidden_dim, num_layers = n_lyrs, dropout = do)
        self.fc1 = nn.Linear(in_features = hidden_dim, out_features = int(hidden_dim / 2))
        self.act1 = nn.ReLU(inplace = True)
        # self.bn1 = nn.BatchNorm1d(num_features = int(hidden_dim / 2))

        self.estimator = nn.Linear(in_features = int(hidden_dim / 2), out_features = 1)
        
    
    def init_hiddenState(self, bs):
        """Initialize the hidden state of RNN to all zeros
        Args:
            bs (int): Batch size during training
        """
        return torch.zeros(self.n_layers, bs, self.hidden_dim)

    def forward(self, input_):
        """Define the forward propogation logic here
        Args:
            input (Tensor): A 3-dimensional float tensor containing parameters
        """
        bs = input_.shape[1]
        hidden_state = self.init_hiddenState(bs).to(self.device)
        cell_state = hidden_state
        
        out, _ = self.rnn(input_, (hidden_state, cell_state))

        out = out.contiguous().view(-1, self.hidden_dim)
        # out = self.act1(self.bn1(self.fc1(out)))
        out = self.act1(self.fc1(out))
        out = self.estimator(out)
        
        return out
    
    def predict(self, input_):
        """Makes prediction for the set of inputs provided and returns the same
        Args:
            input (torch.Tensor): A tensor of inputs
        """
        with torch.no_grad():
            predictions = self.forward(input_)
        
        return predictions