| 26 | |
| 27 | class NavieComplexLSTM(nn.Module): |
| 28 | def __init__(self, input_size, hidden_size, projection_dim=None, bidirectional=False, batch_first=False): |
| 29 | super(NavieComplexLSTM, self).__init__() |
| 30 | |
| 31 | self.input_dim = input_size//2 |
| 32 | self.rnn_units = hidden_size//2 |
| 33 | self.real_lstm = nn.LSTM(self.input_dim, self.rnn_units, num_layers=1, bidirectional=bidirectional, batch_first=False) |
| 34 | self.imag_lstm = nn.LSTM(self.input_dim, self.rnn_units, num_layers=1, bidirectional=bidirectional, batch_first=False) |
| 35 | if bidirectional: |
| 36 | bidirectional=2 |
| 37 | else: |
| 38 | bidirectional=1 |
| 39 | if projection_dim is not None: |
| 40 | self.projection_dim = projection_dim//2 |
| 41 | self.r_trans = nn.Linear(self.rnn_units*bidirectional, self.projection_dim) |
| 42 | self.i_trans = nn.Linear(self.rnn_units*bidirectional, self.projection_dim) |
| 43 | else: |
| 44 | self.projection_dim = None |
| 45 | |
| 46 | def forward(self, inputs): |
| 47 | if isinstance(inputs,list): |