Construct the sub-networks. Overview of the model: (1) A :py:class:`NetworkGraves` that handles all the rendering, or P(pt | zt, x0, ..., x_{t-1}, cs) (2) A multi-head dot-product attention that focuses and extracts info from fs for posterior computation.
(self)
| 232 | self._construct_networks() |
| 233 | |
| 234 | def _construct_networks(self): |
| 235 | """Construct the sub-networks. |
| 236 | |
| 237 | Overview of the model: |
| 238 | (1) A :py:class:`NetworkGraves` that handles all the rendering, or P(pt | zt, x0, ..., x_{t-1}, cs) |
| 239 | (2) A multi-head dot-product attention that focuses and extracts info from fs for posterior computation. |
| 240 | (3) A fully connected net to model the posterior q(zt | fs, cs, x0, ..., x_{t-1}) |
| 241 | (4) A fully connected net to model the prior p(zt | cs, x0, ..., x_{t-1}) |
| 242 | """ |
| 243 | |
| 244 | # alex graves model as the renderer |
| 245 | self.net_graves = NetworkGraves(self.param_graves) |
| 246 | |
| 247 | # copy essential hyper-parameters from param_graves |
| 248 | self.dim_x = self.net_graves.dim_x |
| 249 | self.dim_c = self.net_graves.dim_c |
| 250 | self.dim_z = self.net_graves.dim_z |
| 251 | self.dim_attn_rnn_h = self.net_graves.dim_attn_rnn_h |
| 252 | self.dim_decode_rnn_h = self.net_graves.dim_decode_rnn_h |
| 253 | |
| 254 | # query net (mix attn_hs and attn_cs to the query) |
| 255 | # input: |
| 256 | # attn_hs |
| 257 | # attn_cs |
| 258 | # output: |
| 259 | # query |
| 260 | dim_mix_layer_input = self.dim_attn_rnn_h + self.dim_c |
| 261 | self.net_query = StackedLinearLayers( |
| 262 | num_layers=self.num_query_layers, |
| 263 | dim_input=dim_mix_layer_input, |
| 264 | dim_output=self.dim_query, |
| 265 | dim_features=self.dim_query_layers, |
| 266 | nonlinearity=self.query_nonlinearity, |
| 267 | add_norm_layer=self.query_add_norm, |
| 268 | norm_fun=nn.LayerNorm, |
| 269 | dropout_prob=self.query_dropout_prob, |
| 270 | output_add_nonlinearity=True, # added because multi-head attention has a linear transform at input |
| 271 | ) |
| 272 | if self.flag_bckwrd_compatible: |
| 273 | self.net_query = nn.Sequential(self.net_query) |
| 274 | |
| 275 | # multi-head dot-product attention |
| 276 | # key: fs |
| 277 | # value: fs |
| 278 | # query: net_query(attn_ht, attn_cs) |
| 279 | self.net_attn = nn.MultiheadAttention( |
| 280 | embed_dim=self.dim_query, |
| 281 | num_heads=self.num_attn_heads, |
| 282 | dropout=self.attn_dropout_prob, |
| 283 | kdim=self.dim_f, |
| 284 | vdim=self.dim_f, |
| 285 | ) |
| 286 | |
| 287 | # posterior |
| 288 | # input: attended values (seq_len, batch, dim_query) |
| 289 | # output: posterior distribution parameters (mean, std), assuming independent Gaussian |
| 290 | if self.posterior_nonlinearity == "leaky_relu": |
| 291 | nonlinearity_fun = lambda: nn.LeakyReLU(negative_slope=0.01, inplace=False) |
no test coverage detected