Attention with mixture of gaussian windows. Ref: https://arxiv.org/abs/1308.0850 page 26
| 12 | |
| 13 | |
| 14 | class GaussianSlidingWindows(nn.Module): |
| 15 | """ |
| 16 | Attention with mixture of gaussian windows. |
| 17 | |
| 18 | Ref: https://arxiv.org/abs/1308.0850 page 26 |
| 19 | """ |
| 20 | |
| 21 | def __init__( |
| 22 | self, |
| 23 | num_mixtures: int, |
| 24 | dim_input: int, |
| 25 | num_layers: int, |
| 26 | dim_features: T.Union[int, T.Sequence[int]], |
| 27 | pos_fun: str = "softplus", |
| 28 | normalized_weights=False, |
| 29 | ): |
| 30 | """ |
| 31 | Given input, the layer computes a mixture of 1D gaussians as the attention weights. |
| 32 | |
| 33 | Args: |
| 34 | num_mixtures: |
| 35 | how many gaussian to use |
| 36 | dim_input: |
| 37 | dimension of the input (ex: hidden state of bottom lstm) |
| 38 | num_layers: |
| 39 | number of layers |
| 40 | dim_features: |
| 41 | feature dimension of the linear layers. |
| 42 | Can be an integer so every layer shares the same dimension |
| 43 | or a list of (num_layers-1) integers, one for each layer except the last layer. |
| 44 | pos_fun: |
| 45 | name of the function to learn positive variance. |
| 46 | Choices: [`exp` | `softplus`] |
| 47 | normalized_weights: |
| 48 | whether to add a softmax to gaussian weights |
| 49 | """ |
| 50 | |
| 51 | super().__init__() |
| 52 | self.num_mixtures = num_mixtures |
| 53 | self.dim_input = dim_input |
| 54 | self.dim_features = dim_features |
| 55 | self.num_layers = num_layers |
| 56 | self.dim_linear_output = self.num_mixtures * 3 # mean, std, and weight for each gaussian |
| 57 | if pos_fun == "exp": |
| 58 | self.make_pos_fun = lambda x: torch.exp(x) |
| 59 | elif pos_fun == "softplus": |
| 60 | self.softplus = nn.Softplus() |
| 61 | self.make_pos_fun = lambda x: self.softplus(x) |
| 62 | else: |
| 63 | raise ValueError("wrong") |
| 64 | |
| 65 | self.normalized_weights = normalized_weights |
| 66 | if self.normalized_weights: |
| 67 | self.log_softmax = nn.LogSoftmax(dim=2) |
| 68 | |
| 69 | self.main = StackedLinearLayers( |
| 70 | num_layers=self.num_layers, |
| 71 | dim_input=self.dim_input, |