Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| r""" | |
| Layer-Wise Attention Mechanism | |
| ================================ | |
| Computes a parameterised scalar mixture of N tensors, `mixture = gamma * sum(s_k * tensor_k)` | |
| where `s = softmax(w)`, with `w` and `gamma` scalar parameters. | |
| If `do_layer_norm=True` then apply layer normalization to each tensor before weighting. | |
| If `dropout > 0`, then for each scalar weight, adjust its softmax weight mass to 0 with | |
| the dropout probability (i.e., setting the unnormalized weight to -inf). This effectively | |
| should redistribute dropped probability mass to all other weights. | |
| Original implementation: | |
| - https://github.com/Hyperparticle/udify | |
| """ | |
| from typing import List | |
| import torch | |
| from torch.nn import Parameter, ParameterList | |
| class ScalarMixWithDropout(torch.nn.Module): | |
| def __init__( | |
| self, | |
| mixture_size: int, | |
| do_layer_norm: bool = False, | |
| initial_scalar_parameters: list = None, | |
| trainable: bool = True, | |
| dropout: float = None, | |
| dropout_value: float = -1e20, | |
| ) -> None: | |
| super(ScalarMixWithDropout, self).__init__() | |
| self.mixture_size = mixture_size | |
| self.do_layer_norm = do_layer_norm | |
| self.dropout = dropout | |
| if initial_scalar_parameters is None: | |
| initial_scalar_parameters = [0.0] * mixture_size | |
| elif len(initial_scalar_parameters) != mixture_size: | |
| raise Exception( | |
| "Length of initial_scalar_parameters {} differs \ | |
| from mixture_size {}".format( | |
| initial_scalar_parameters, mixture_size | |
| ) | |
| ) | |
| self.scalar_parameters = ParameterList( | |
| [ | |
| Parameter( | |
| torch.FloatTensor([initial_scalar_parameters[i]]), | |
| requires_grad=trainable, | |
| ) | |
| for i in range(mixture_size) | |
| ] | |
| ) | |
| self.gamma = Parameter(torch.FloatTensor([1.0]), requires_grad=trainable) | |
| if self.dropout: | |
| dropout_mask = torch.zeros(len(self.scalar_parameters)) | |
| dropout_fill = torch.empty(len(self.scalar_parameters)).fill_(dropout_value) | |
| self.register_buffer("dropout_mask", dropout_mask) | |
| self.register_buffer("dropout_fill", dropout_fill) | |
| def forward( | |
| self, | |
| tensors: List[torch.Tensor], # pylint: disable=arguments-differ | |
| mask: torch.Tensor = None, | |
| ) -> torch.Tensor: | |
| """ | |
| Compute a weighted average of the `tensors`. The input tensors an be any shape | |
| with at least two dimensions, but must all be the same shape. | |
| When `do_layer_norm=True`, the `mask` is required input. If the `tensors` are | |
| dimensioned `(dim_0, ..., dim_{n-1}, dim_n)`, then the `mask` is dimensioned | |
| `(dim_0, ..., dim_{n-1})`, as in the typical case with `tensors` of shape | |
| `(batch_size, timesteps, dim)` and `mask` of shape `(batch_size, timesteps)`. | |
| When `do_layer_norm=False` the `mask` is ignored. | |
| """ | |
| if len(tensors) != self.mixture_size: | |
| raise Exception( | |
| "{} tensors were passed, but the module was initialized to \ | |
| mix {} tensors.".format( | |
| len(tensors), self.mixture_size | |
| ) | |
| ) | |
| def _do_layer_norm(tensor, broadcast_mask, num_elements_not_masked): | |
| tensor_masked = tensor * broadcast_mask | |
| mean = torch.sum(tensor_masked) / num_elements_not_masked | |
| variance = ( | |
| torch.sum(((tensor_masked - mean) * broadcast_mask) ** 2) | |
| / num_elements_not_masked | |
| ) | |
| return (tensor - mean) / torch.sqrt(variance + 1e-12) | |
| weights = torch.cat([parameter for parameter in self.scalar_parameters]) | |
| if self.training and self.dropout: | |
| weights = torch.where( | |
| self.dropout_mask.uniform_() > self.dropout, weights, self.dropout_fill | |
| ) | |
| normed_weights = torch.nn.functional.softmax(weights, dim=0) | |
| normed_weights = torch.split(normed_weights, split_size_or_sections=1) | |
| if not self.do_layer_norm: | |
| pieces = [] | |
| for weight, tensor in zip(normed_weights, tensors): | |
| pieces.append(weight * tensor) | |
| return self.gamma * sum(pieces) | |
| else: | |
| mask_float = mask.float() | |
| broadcast_mask = mask_float.unsqueeze(-1) | |
| input_dim = tensors[0].size(-1) | |
| num_elements_not_masked = torch.sum(mask_float) * input_dim | |
| pieces = [] | |
| for weight, tensor in zip(normed_weights, tensors): | |
| pieces.append( | |
| weight | |
| * _do_layer_norm(tensor, broadcast_mask, num_elements_not_masked) | |
| ) | |
| return self.gamma * sum(pieces) | |