Adding layers to nn.RNN

This is my code:

class gwRNN(nn.Module):
  #declare a few variables with initialisation function:
    def __init__(self, batch_size, n_steps, n_inputs, n_neurons, n_real_outputs, num_layers):
        super(gwRNN, self).__init__()
        self.batch_size = batch_size
        self.n_steps = n_steps
        self.n_inputs = n_inputs
        self.n_neurons = n_neurons
        self.n_real_outputs = n_real_outputs
        self.num_layers = num_layers

        # declare basic RNN layer
        self.basic_rnn = nn.RNN(self.n_inputs, self.n_neurons, num_layers, nonlinearity = 'relu', dropout=0)
        # declare fully-connected layer 
        self.FC = nn.Linear(self.n_neurons, self.n_real_outputs, bias= False) #nn.linear Applies a linear transformation to the incoming data
    def init_hidden(self,): # returns hidden state with zero values
        # (num_layers, batch_size, n_neurons)
        return (torch.zeros(self.num_layers, self.batch_size, self.n_neurons))
    def forward(self, X):
        # transforms X to dimensions: n_steps X batch_size X n_inputs
        X = X.permute(1, 0, 2) 

        # The data flows through the RNN layer and then through the fully-connected layer.        
        self.batch_size = X.size(1) # size of 2nd element of X
        self.hidden = self.init_hidden()
        lstm_out, self.hidden = self.basic_rnn(X, self.hidden)   #SEE ABOVE for self.basic_rnn and nn.RNN !!    
        out = self.FC(self.hidden) # output represent the log probabilities of the model.  
        return out.view(-1, self.n_real_outputs) # batch_size X n_output

I’d like to add stack rnns by changing num_layers, but when I do this my code doesn’t work as it returns out with shape (n_layers*batch size) X n_output instead of batch_size X n_output

You permuted X to make batch the second dimension. Going by that, the self.hidden has to be permuted back to make batch the first dimension before feeding to self.FC.

1 Like

just slice it to the last layer: self.hidden[-1]

1 Like