def test_res_rnn(rnn_type, dropout, bidir): n_units, n_layers = 20, 3 model = blocks.StackedResidualRNN(rnn_type, n_units, n_layers=n_layers, dropout=dropout, bidirectional=bidir) batch, n_frames = 2, 78 inp = torch.randn(batch, n_frames, n_units) out = model(inp) assert out.shape == (batch, n_frames, n_units)
def __init__(self, input_size, hidden_size, output_size=None, rnn_type='gru', n_layers=3, dropout=0.3): super(SimpleModel, self).__init__() self.input_size = input_size self.hidden_size = hidden_size output_size = input_size if output_size is None else output_size self.output_size = output_size self.in_proj_layer = nn.Linear(input_size, hidden_size) self.residual_rec = blocks.StackedResidualRNN(rnn_type, hidden_size, n_layers=n_layers, dropout=dropout) self.out_proj_layer = nn.Linear(hidden_size, output_size)