Args: x: (seq_len, batch_size, dim_input) h0: hidden state of the lstm. See :py:class:`nn.LSTM`. Returns: y: (seq_len, batch_size, dim) or (batch_size, dim) depending on `output_mode` h:
(self, x, h0=None, valid_seq_lens=None)
| 253 | ) |
| 254 | |
| 255 | def forward(self, x, h0=None, valid_seq_lens=None): |
| 256 | """ |
| 257 | Args: |
| 258 | x: (seq_len, batch_size, dim_input) |
| 259 | h0: |
| 260 | hidden state of the lstm. See :py:class:`nn.LSTM`. |
| 261 | Returns: |
| 262 | y: |
| 263 | (seq_len, batch_size, dim) or (batch_size, dim) depending on `output_mode` |
| 264 | h: |
| 265 | final hidden state of the lstm. See :py:class:`nn.LSTM`. |
| 266 | |
| 267 | Note when we output_final_only, i.e., y is the shape of (batch_size, dim_output) |
| 268 | if valid_seq_lens is None: |
| 269 | return the output at seq_len-1 (final time step) |
| 270 | else: |
| 271 | return the output of valid_seq_lens-1 |
| 272 | """ |
| 273 | seq_len = x.size(0) |
| 274 | if valid_seq_lens is not None: |
| 275 | if isinstance(valid_seq_lens, torch.Tensor): |
| 276 | valid_seq_lens = valid_seq_lens.detach().cpu() |
| 277 | x = nn.utils.rnn.pack_padded_sequence(x, valid_seq_lens, batch_first=False, enforce_sorted=False) |
| 278 | |
| 279 | # run rnn on x |
| 280 | out, h = self.rnn(x, h0) |
| 281 | |
| 282 | if isinstance(out, torch.nn.utils.rnn.PackedSequence): |
| 283 | assert valid_seq_lens is not None |
| 284 | out, _ = nn.utils.rnn.pad_packed_sequence( |
| 285 | out, |
| 286 | batch_first=False, |
| 287 | padding_value=0.0, |
| 288 | total_length=seq_len, |
| 289 | ) |
| 290 | |
| 291 | if self.output_mode == "last_valid": |
| 292 | assert valid_seq_lens is not None |
| 293 | linear_inputs = torch.zeros(out.size(1), self.rnn_feature_dim, device=out.device) |
| 294 | for b in range(out.size(1)): |
| 295 | linear_inputs[b] = out[valid_seq_lens[b] - 1, b] |
| 296 | elif self.output_mode == "all": |
| 297 | linear_inputs = out |
| 298 | elif self.output_mode == "last": |
| 299 | linear_inputs = out[-1] |
| 300 | elif self.output_mode == "max_valid": |
| 301 | assert valid_seq_lens is not None |
| 302 | linear_inputs = torch.zeros(out.size(1), self.rnn_feature_dim, device=out.device) |
| 303 | for b in range(out.size(1)): |
| 304 | linear_inputs[b], _ = torch.max(out[: valid_seq_lens[b], b], dim=0) |
| 305 | elif self.output_mode == "max": |
| 306 | linear_inputs, _ = torch.max(out, dim=0) |
| 307 | elif self.output_mode == "avg_valid": |
| 308 | assert valid_seq_lens is not None |
| 309 | linear_inputs = torch.zeros(out.size(1), self.rnn_feature_dim, device=out.device) |
| 310 | for b in range(out.size(1)): |
| 311 | linear_inputs[b] = torch.mean(out[: valid_seq_lens[b], b], dim=0) |
| 312 | elif self.output_mode == "avg": |