Convenient helper nn.Module to create a stack of RNNCells. Note that nn.LSTMCell is preferred over nn.LSTM when for loop needs to be called during training.
| 15 | |
| 16 | |
| 17 | class LSTMCellLayers(nn.Module): |
| 18 | """ |
| 19 | Convenient helper nn.Module to create a stack of RNNCells. |
| 20 | |
| 21 | Note that nn.LSTMCell is preferred over nn.LSTM when |
| 22 | for loop needs to be called during training. |
| 23 | """ |
| 24 | |
| 25 | def __init__( |
| 26 | self, |
| 27 | num_layers: int, |
| 28 | dim_input: int, |
| 29 | dim_hidden: T.Union[T.Sequence[int], int], |
| 30 | bias: bool = True, |
| 31 | append_inputs: bool = False, |
| 32 | lstm_cell_fn: T.Callable = nn.LSTMCell, |
| 33 | dropout_prob: float = 0, |
| 34 | ): |
| 35 | """ |
| 36 | A multi-layer RNN composed of LSTMCells (or the provided lstm_cell_fn). |
| 37 | |
| 38 | Args: |
| 39 | num_layers: |
| 40 | Number of RNN layers to create. |
| 41 | dim_input: |
| 42 | Feature dimension of the input. |
| 43 | dim_hidden: |
| 44 | Dimension of the hidden states. If an integer is provided, it will be used for all layers. |
| 45 | Otherwise, provide a list of integer, one for each layer. |
| 46 | bias: |
| 47 | Whether to learn bias at each layer. |
| 48 | append_inputs: |
| 49 | Whether to concatenate all previous layers' inputs to every layer's input. |
| 50 | Note that the input to i-th layer is the output hidden state of (i-1)-th layer (plus the |
| 51 | input to all previous layers if append_inputs is True.) |
| 52 | lstm_cell_fn: |
| 53 | RNNCell function to use. Can be one of |
| 54 | :py:func:`torch.nn.RNNCell`, |
| 55 | :py:func:`torch.nn.LSTMCell`, |
| 56 | :py:func:`torch.nn.GRUCell`. |
| 57 | dropout_prob: |
| 58 | Dropout probability on the hidden states. If non-zero, introduces a Dropout layer on the outputs of |
| 59 | each LSTM layer except the last layer. |
| 60 | |
| 61 | """ |
| 62 | super().__init__() |
| 63 | self.num_layers = num_layers |
| 64 | self.dim_input = dim_input |
| 65 | self.dim_hidden = dim_hidden |
| 66 | self.lstm_cell_fn = lstm_cell_fn |
| 67 | self.dropout_prob = dropout_prob |
| 68 | if isinstance(self.dim_hidden, int): |
| 69 | self.dim_hidden = [self.dim_hidden for _ in range(self.num_layers)] |
| 70 | elif len(self.dim_hidden) == 1: |
| 71 | self.dim_hidden = [self.dim_hidden[0] for _ in range(self.num_layers)] |
| 72 | self.append_inputs = append_inputs |
| 73 | self.cell_list = nn.ModuleList() |
| 74 | din = dim_input |