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)