A LSTM with additional linear layers at the top. Args: dim_input: input dimension of the lstm dim_output: output dimension of the linear layer num_rnn_layers: int number of lstm layers r
(
self,
dim_input: int,
dim_output: int,
num_rnn_layers: int,
rnn_feature_dim: int,
num_linear_layers: int,
dim_linear_features: int,
linear_add_layer_norm: bool = False,
output_mode: str = "last_valid",
dropout: float = 0,
bidirectional: bool = False,
)
| 162 | |
| 163 | class LSTMLinear(nn.Module): |
| 164 | def __init__( |
| 165 | self, |
| 166 | dim_input: int, |
| 167 | dim_output: int, |
| 168 | num_rnn_layers: int, |
| 169 | rnn_feature_dim: int, |
| 170 | num_linear_layers: int, |
| 171 | dim_linear_features: int, |
| 172 | linear_add_layer_norm: bool = False, |
| 173 | output_mode: str = "last_valid", |
| 174 | dropout: float = 0, |
| 175 | bidirectional: bool = False, |
| 176 | ): |
| 177 | """ |
| 178 | A LSTM with additional linear layers at the top. |
| 179 | |
| 180 | Args: |
| 181 | dim_input: |
| 182 | input dimension of the lstm |
| 183 | dim_output: |
| 184 | output dimension of the linear layer |
| 185 | num_rnn_layers: int |
| 186 | number of lstm layers |
| 187 | rnn_feature_dim: int |
| 188 | feature dimension of the lstm. See :py:class:`nn.LSTM`. |
| 189 | num_linear_layers: int |
| 190 | number of linear layers |
| 191 | dim_linear_features: |
| 192 | list of feature dimensions of the linear layers. length: num_linear_layers-1 |
| 193 | linear_add_layer_norm: |
| 194 | whether to add layer norm in the linear layers. |
| 195 | output_mode: |
| 196 | ["all" | "last_valid" | "last" | "max_valid" | "max" | "avg_valid" | "avg"] |
| 197 | dropout: |
| 198 | dropout probability on both lstm and linear layers |
| 199 | """ |
| 200 | super().__init__() |
| 201 | |
| 202 | self.dim_input = dim_input |
| 203 | self.dim_output = dim_output |
| 204 | self.num_rnn_layers = num_rnn_layers |
| 205 | self.rnn_feature_dim = rnn_feature_dim |
| 206 | self.num_linear_layers = num_linear_layers |
| 207 | self.dim_linear_features = dim_linear_features |
| 208 | self.output_mode = output_mode |
| 209 | self.linear_add_layer_norm = linear_add_layer_norm |
| 210 | self.bidirectional = bidirectional |
| 211 | assert self.output_mode in { |
| 212 | "all", |
| 213 | "last_valid", |
| 214 | "last", |
| 215 | "max_valid", |
| 216 | "max", |
| 217 | "avg_valid", |
| 218 | "avg", |
| 219 | } |
| 220 | |
| 221 | self.rnn = nn.LSTM( |
no test coverage detected