| 25 | return torch.cat([real,imag],self.complex_axis) |
| 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): |
| 48 | real, imag = inputs |
| 49 | elif isinstance(inputs, torch.Tensor): |
| 50 | real, imag = torch.chunk(inputs,-1) |
| 51 | r2r_out = self.real_lstm(real)[0] |
| 52 | r2i_out = self.imag_lstm(real)[0] |
| 53 | i2r_out = self.real_lstm(imag)[0] |
| 54 | i2i_out = self.imag_lstm(imag)[0] |
| 55 | real_out = r2r_out - i2i_out |
| 56 | imag_out = i2r_out + r2i_out |
| 57 | if self.projection_dim is not None: |
| 58 | real_out = self.r_trans(real_out) |
| 59 | imag_out = self.i_trans(imag_out) |
| 60 | #print(real_out.shape,imag_out.shape) |
| 61 | return [real_out, imag_out] |
| 62 | |
| 63 | def flatten_parameters(self): |
| 64 | self.imag_lstm.flatten_parameters() |
| 65 | self.real_lstm.flatten_parameters() |
| 66 | |
| 67 | def complex_cat(inputs, axis): |
| 68 |
nothing calls this directly
no outgoing calls
no test coverage detected