Construct the sub-networks. Overview of the model: attn_rnn (lstm_cell_layers): input: x_{t-1}, last_attn_context output: ht attn_layer (gaussian window attention): input: ht, output: attn_weights, attn_context,
(self)
| 219 | self._construct_networks() |
| 220 | |
| 221 | def _construct_networks(self): |
| 222 | """Construct the sub-networks. |
| 223 | |
| 224 | Overview of the model: |
| 225 | attn_rnn (lstm_cell_layers): |
| 226 | input: x_{t-1}, last_attn_context |
| 227 | output: ht |
| 228 | attn_layer (gaussian window attention): |
| 229 | input: ht, |
| 230 | output: attn_weights, attn_context, attn_delta_means |
| 231 | decode_rnn (LSTM): |
| 232 | input: attn_context, ht, x_{t-1}, (zt, if provided) |
| 233 | output: raw_output |
| 234 | output_layer: |
| 235 | input: raw_output |
| 236 | output: pt |
| 237 | """ |
| 238 | |
| 239 | ## attn_rnn |
| 240 | # input: |
| 241 | # (1) x_{t-1} (batch, dim_x) |
| 242 | # (2) last_attn_context (batch, dim_c) |
| 243 | dim_attn_rnn_input = self.dim_x + self.dim_c |
| 244 | self.attn_rnn = LSTMCellLayers( |
| 245 | num_layers=self.num_attn_rnn_layers, |
| 246 | dim_input=dim_attn_rnn_input, |
| 247 | dim_hidden=self.dim_attn_rnn, |
| 248 | ) |
| 249 | self.dim_attn_rnn_h = self.attn_rnn.dim_hidden[-1] |
| 250 | |
| 251 | ## attn_layer |
| 252 | # input: attn_rnn_h |
| 253 | # output: attn_weights, attn_context, attn_delta_means |
| 254 | dim_attn_layer_input = self.dim_attn_rnn_h |
| 255 | self.attn_layer = GaussianSlidingWindows( |
| 256 | num_mixtures=self.num_window_attn_mixtures, |
| 257 | dim_input=dim_attn_layer_input, |
| 258 | num_layers=self.num_window_attn_layers, |
| 259 | dim_features=self.dim_window_attn_layers, |
| 260 | pos_fun="softplus", |
| 261 | ) |
| 262 | |
| 263 | ## decode rnn |
| 264 | # input: |
| 265 | # (1) attn_context (batch, dim_c) |
| 266 | # (2) h_attn_rnn |
| 267 | # (3) x_{t-1} |
| 268 | # (4) z (batch, dim_z), if provided |
| 269 | dim_decode_rnn_input = self.dim_c + self.dim_attn_rnn_h + self.dim_x |
| 270 | if self.feed_z_to_decode_rnn: |
| 271 | dim_decode_rnn_input += self.dim_z |
| 272 | |
| 273 | self.decode_rnn = nn.LSTM( |
| 274 | input_size=dim_decode_rnn_input, |
| 275 | hidden_size=self.dim_decode_rnn, |
| 276 | num_layers=self.num_decode_rnn_layers, |
| 277 | batch_first=False, |
| 278 | bidirectional=False, |
no test coverage detected