I’m trying to implement an RNN with Hierarchical Subsampling as shown here [chapter 9.0]

image920×666 82.6 KB

Is my understanding of how to implement this correct? I feel that this is too simple
import torch
import torch.nn as nn

class subsampleRNN(nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.lstm1 = nn.LSTM(512, 256)
        self.lstm2 = nn.LSTM(256, 128)
        self.lstm3 = nn.LSTM(128, 64)
        
        self.relu = nn.ReLU()
        
         
    def forward(self, x):
        x, h0 = self.lstm1(x)
        x = self.relu(x)
        x, h0 = self.lstm2(x)
        x = self.relu(x)
        x, h0 = self.lstm3(x)
        
        return x
    
    
model = subsampleRNN()
input = torch.rand(1,1,512)
model(input).shape