| |
| |
| |
| |
|
|
| import functools |
| import operator |
|
|
| import torch |
| import torch.nn.functional as F |
| from fairseq.modules.fairseq_dropout import FairseqDropout |
| from fairseq.modules.quant_noise import quant_noise |
| from torch import nn |
|
|
|
|
| class TiedLinear(nn.Module): |
| def __init__(self, weight, transpose): |
| super().__init__() |
| self.weight = weight |
| self.transpose = transpose |
|
|
| def forward(self, input): |
| return F.linear(input, self.weight.t() if self.transpose else self.weight) |
|
|
|
|
| class TiedHeadModule(nn.Module): |
| def __init__(self, weights, input_dim, num_classes, q_noise, qn_block_size): |
| super().__init__() |
| tied_emb, _ = weights |
| self.num_words, emb_dim = tied_emb.size() |
|
|
| self.word_proj = quant_noise( |
| TiedLinear(tied_emb, transpose=False), q_noise, qn_block_size |
| ) |
| if input_dim != emb_dim: |
| self.word_proj = nn.Sequential( |
| quant_noise( |
| nn.Linear(input_dim, emb_dim, bias=False), q_noise, qn_block_size |
| ), |
| self.word_proj, |
| ) |
|
|
| self.class_proj = quant_noise( |
| nn.Linear(input_dim, num_classes, bias=False), q_noise, qn_block_size |
| ) |
| self.out_dim = self.num_words + num_classes |
|
|
| self.register_buffer("_float_tensor", torch.FloatTensor(1)) |
|
|
| def forward(self, input): |
| inp_sz = functools.reduce(operator.mul, input.shape[:-1], 1) |
| out = self._float_tensor.new(inp_sz, self.out_dim) |
| out[:, : self.num_words] = self.word_proj(input.view(inp_sz, -1)) |
| out[:, self.num_words :] = self.class_proj(input.view(inp_sz, -1)) |
| return out |
|
|
|
|
| class AdaptiveSoftmax(nn.Module): |
| """ |
| This is an implementation of the efficient softmax approximation for |
| graphical processing units (GPU), described in the paper "Efficient softmax |
| approximation for GPUs" (http://arxiv.org/abs/1609.04309). |
| """ |
|
|
| def __init__( |
| self, |
| vocab_size, |
| input_dim, |
| cutoff, |
| dropout, |
| factor=4.0, |
| adaptive_inputs=None, |
| tie_proj=False, |
| q_noise=0, |
| qn_block_size=8, |
| ): |
| super().__init__() |
|
|
| if vocab_size > cutoff[-1]: |
| cutoff = cutoff + [vocab_size] |
| else: |
| assert ( |
| vocab_size == cutoff[-1] |
| ), "cannot specify cutoff larger than vocab size" |
|
|
| output_dim = cutoff[0] + len(cutoff) - 1 |
|
|
| self.vocab_size = vocab_size |
| self.cutoff = cutoff |
| self.dropout_module = FairseqDropout( |
| dropout, module_name=self.__class__.__name__ |
| ) |
| self.input_dim = input_dim |
| self.factor = factor |
| self.q_noise = q_noise |
| self.qn_block_size = qn_block_size |
|
|
| self.lsm = nn.LogSoftmax(dim=1) |
|
|
| if adaptive_inputs is not None: |
| self.head = TiedHeadModule( |
| adaptive_inputs.weights_for_band(0), |
| input_dim, |
| len(cutoff) - 1, |
| self.q_noise, |
| self.qn_block_size, |
| ) |
| else: |
| self.head = quant_noise( |
| nn.Linear(input_dim, output_dim, bias=False), |
| self.q_noise, |
| self.qn_block_size, |
| ) |
|
|
| self._make_tail(adaptive_inputs, tie_proj) |
|
|
| def init_weights(m): |
| if ( |
| hasattr(m, "weight") |
| and not isinstance(m, TiedLinear) |
| and not isinstance(m, TiedHeadModule) |
| ): |
| nn.init.xavier_uniform_(m.weight) |
|
|
| self.apply(init_weights) |
|
|
| self.register_buffer("version", torch.LongTensor([1])) |
|
|
| def _make_tail(self, adaptive_inputs=None, tie_proj=False): |
| self.tail = nn.ModuleList() |
| for i in range(len(self.cutoff) - 1): |
| dim = int(self.input_dim // self.factor ** (i + 1)) |
|
|
| tied_emb, tied_proj = ( |
| adaptive_inputs.weights_for_band(i + 1) |
| if adaptive_inputs is not None |
| else (None, None) |
| ) |
|
|
| if tied_proj is not None: |
| if tie_proj: |
| proj = quant_noise( |
| TiedLinear(tied_proj, transpose=True), |
| self.q_noise, |
| self.qn_block_size, |
| ) |
| else: |
| proj = quant_noise( |
| nn.Linear(tied_proj.size(0), tied_proj.size(1), bias=False), |
| self.q_noise, |
| self.qn_block_size, |
| ) |
| else: |
| proj = quant_noise( |
| nn.Linear(self.input_dim, dim, bias=False), |
| self.q_noise, |
| self.qn_block_size, |
| ) |
|
|
| if tied_emb is None: |
| out_proj = nn.Linear( |
| dim, self.cutoff[i + 1] - self.cutoff[i], bias=False |
| ) |
| else: |
| out_proj = TiedLinear(tied_emb, transpose=False) |
|
|
| m = nn.Sequential( |
| proj, |
| nn.Dropout(self.dropout_module.p), |
| quant_noise(out_proj, self.q_noise, self.qn_block_size), |
| ) |
|
|
| self.tail.append(m) |
|
|
| def upgrade_state_dict_named(self, state_dict, name): |
| version_name = name + ".version" |
| if version_name not in state_dict: |
| raise Exception("This version of the model is no longer supported") |
|
|
| def adapt_target(self, target): |
| """ |
| In order to be efficient, the AdaptiveSoftMax does not compute the |
| scores for all the word of the vocabulary for all the examples. It is |
| thus necessary to call the method adapt_target of the AdaptiveSoftMax |
| layer inside each forward pass. |
| """ |
|
|
| target = target.view(-1) |
| new_target = [target.clone()] |
| target_idxs = [] |
|
|
| for i in range(len(self.cutoff) - 1): |
| mask = target.ge(self.cutoff[i]).mul(target.lt(self.cutoff[i + 1])) |
| new_target[0][mask] = self.cutoff[0] + i |
|
|
| if mask.any(): |
| target_idxs.append(mask.nonzero(as_tuple=False).squeeze(1)) |
| new_target.append(target[mask].add(-self.cutoff[i])) |
| else: |
| target_idxs.append(None) |
| new_target.append(None) |
|
|
| return new_target, target_idxs |
|
|
| def forward(self, input, target): |
| """ |
| Args: |
| input: (b x t x d) |
| target: (b x t) |
| Returns: |
| 2 lists: output for each cutoff section and new targets by cut off |
| """ |
|
|
| input = input.contiguous().view(-1, input.size(-1)) |
| input = self.dropout_module(input) |
|
|
| new_target, target_idxs = self.adapt_target(target) |
| output = [self.head(input)] |
|
|
| for i in range(len(target_idxs)): |
| if target_idxs[i] is not None: |
| output.append(self.tail[i](input.index_select(0, target_idxs[i]))) |
| else: |
| output.append(None) |
|
|
| return output, new_target |
|
|
| def get_log_prob(self, input, target): |
| """ |
| Computes the log probabilities for all the words of the vocabulary, |
| given a 2D tensor of hidden vectors. |
| """ |
|
|
| bsz, length, dim = input.size() |
| input = input.contiguous().view(-1, dim) |
|
|
| if target is not None: |
| _, target_idxs = self.adapt_target(target) |
| else: |
| target_idxs = None |
|
|
| head_y = self.head(input) |
| log_probs = head_y.new_zeros(input.size(0), self.vocab_size) |
|
|
| head_sz = self.cutoff[0] + len(self.tail) |
| log_probs[:, :head_sz] = self.lsm(head_y) |
| tail_priors = log_probs[:, self.cutoff[0] : head_sz].clone() |
|
|
| for i in range(len(self.tail)): |
| start = self.cutoff[i] |
| end = self.cutoff[i + 1] |
|
|
| if target_idxs is None: |
| tail_out = log_probs[:, start:end] |
| tail_out.copy_(self.tail[i](input)) |
| log_probs[:, start:end] = self.lsm(tail_out).add_( |
| tail_priors[:, i, None] |
| ) |
| elif target_idxs[i] is not None: |
| idxs = target_idxs[i] |
| tail_out = log_probs[idxs, start:end] |
| tail_out.copy_(self.tail[i](input[idxs])) |
| log_probs[idxs, start:end] = self.lsm(tail_out).add_( |
| tail_priors[idxs, i, None] |
| ) |
|
|
| log_probs = log_probs.view(bsz, length, -1) |
| return log_probs |
|
|
