Loading fairseq/models/fairseq_incremental_decoder.py +5 −84 Changes for fairseq/models/fairseq_incremental_decoder.py: 5 added lines, 84 removed lines. Original line number Diff line number Diff line Loading @@ -13,99 +13,20 @@ class FairseqIncrementalDecoder(FairseqDecoder): def __init__(self, dictionary): super().__init__(dictionary) self._is_incremental_eval = False self._incremental_state = {} def forward(self, prev_output_tokens, encoder_out): if self._is_incremental_eval: def forward(self, prev_output_tokens, encoder_out, incremental_state=None): raise NotImplementedError else: raise NotImplementedError def incremental_inference(self): """Context manager for incremental inference. This provides an optimized forward pass for incremental inference (i.e., it predicts one time step at a time). If the input order changes between time steps, call reorder_incremental_state to update the relevant buffers. To generate a fresh sequence, first call clear_incremental_state. Usage: ``` with model.decoder.incremental_inference(): for step in range(maxlen): out, _ = model.decoder(tokens[:, :step], encoder_out) probs = model.get_normalized_probs(out[:, -1, :], log_probs=False) ``` """ class IncrementalInference(object): def __init__(self, decoder): self.decoder = decoder def __enter__(self): self.decoder.incremental_eval(True) def __exit__(self, *args): self.decoder.incremental_eval(False) return IncrementalInference(self) def incremental_eval(self, mode=True): """Sets the decoder and all children in incremental evaluation mode.""" assert self._is_incremental_eval != mode, \ 'incremental_eval already set to mode {}'.format(mode) self._is_incremental_eval = mode if mode: self.clear_incremental_state() def apply_incremental_eval(module): if module != self and hasattr(module, 'incremental_eval'): module.incremental_eval(mode) self.apply(apply_incremental_eval) def get_incremental_state(self, key): """Return cached state or None if not in incremental inference mode.""" if self._is_incremental_eval and key in self._incremental_state: return self._incremental_state[key] return None def set_incremental_state(self, key, value): """Cache state needed for incremental inference mode.""" if self._is_incremental_eval: self._incremental_state[key] = value return value def clear_incremental_state(self): """Clear all state used for incremental generation. **For incremental inference only** This should be called before generating a fresh sequence. beam_size is required if using BeamableMM. """ if self._is_incremental_eval: del self._incremental_state self._incremental_state = {} def apply_clear_incremental_state(module): if module != self and hasattr(module, 'clear_incremental_state'): module.clear_incremental_state() self.apply(apply_clear_incremental_state) def reorder_incremental_state(self, new_order): """Reorder buffered internal state (for incremental generation). **For incremental inference only** def reorder_incremental_state(self, incremental_state, new_order): """Reorder incremental state. This should be called when the order of the input has changed from the previous time step. A typical use case is beam search, where the input order changes between time steps based on the choice of beams. order changes between time steps based on the selection of beams. """ if self._is_incremental_eval: def apply_reorder_incremental_state(module): if module != self and hasattr(module, 'reorder_incremental_state'): module.reorder_incremental_state(new_order) module.reorder_incremental_state(incremental_state, new_order) self.apply(apply_reorder_incremental_state) def set_beam_size(self, beam_size): Loading fairseq/models/fconv.py +27 −27 Changes for fairseq/models/fconv.py: 27 added lines, 27 removed lines. Original line number Diff line number Diff line Loading @@ -10,6 +10,7 @@ import torch import torch.nn as nn import torch.nn.functional as F from fairseq import utils from fairseq.data import LanguagePairDataset from fairseq.modules import BeamableMM, GradMultiply, LearnedPositionalEmbedding, LinearizedConvolution Loading Loading @@ -229,19 +230,13 @@ class FConvDecoder(FairseqIncrementalDecoder): else: self.fc3 = Linear(out_embed_dim, num_embeddings, dropout=dropout) def forward(self, prev_output_tokens, encoder_out): def forward(self, prev_output_tokens, encoder_out, incremental_state=None): # split and transpose encoder outputs encoder_a, encoder_b = self._split_encoder_out(encoder_out) encoder_a, encoder_b = self._split_encoder_out(encoder_out, incremental_state) # embed positions positions = self.embed_positions(prev_output_tokens) if self._is_incremental_eval: # keep only the last token for incremental forward pass prev_output_tokens = prev_output_tokens[:, -1:] # embed tokens and positions x = self.embed_tokens(prev_output_tokens) + positions # embed tokens and combine with positional embeddings x = self._embed_tokens(prev_output_tokens, incremental_state) x += self.embed_positions(prev_output_tokens, incremental_state) x = F.dropout(x, p=self.dropout, training=self.training) target_embedding = x Loading @@ -249,7 +244,7 @@ class FConvDecoder(FairseqIncrementalDecoder): x = self.fc1(x) # B x T x C -> T x B x C x = self._transpose_unless_incremental_eval(x) x = self._transpose_if_training(x, incremental_state) # temporal convolutions avg_attn_scores = None Loading @@ -258,13 +253,14 @@ class FConvDecoder(FairseqIncrementalDecoder): residual = x if proj is None else proj(x) x = F.dropout(x, p=self.dropout, training=self.training) x = conv(x) x = conv(x, incremental_state) if incremental_state is None: x = conv.remove_future_timesteps(x) x = F.glu(x, dim=2) # attention if attention is not None: x = self._transpose_unless_incremental_eval(x) x = self._transpose_if_training(x, incremental_state) x, attn_scores = attention(x, target_embedding, (encoder_a, encoder_b)) attn_scores = attn_scores / num_attn_layers Loading @@ -273,13 +269,13 @@ class FConvDecoder(FairseqIncrementalDecoder): else: avg_attn_scores.add_(attn_scores) x = self._transpose_unless_incremental_eval(x) x = self._transpose_if_training(x, incremental_state) # residual x = (x + residual) * math.sqrt(0.5) # T x B x C -> B x T x C x = self._transpose_unless_incremental_eval(x) x = self._transpose_if_training(x, incremental_state) # project back to size of vocabulary x = self.fc2(x) Loading @@ -288,10 +284,6 @@ class FConvDecoder(FairseqIncrementalDecoder): return x, avg_attn_scores def reorder_incremental_state(self, new_order): """Reorder buffered internal state (for incremental generation).""" super().reorder_incremental_state(new_order) def max_positions(self): """Maximum output length supported by the decoder.""" return self.embed_positions.max_positions() Loading @@ -306,13 +298,19 @@ class FConvDecoder(FairseqIncrementalDecoder): state_dict['decoder.version'] = torch.Tensor([1]) return state_dict def _split_encoder_out(self, encoder_out): def _embed_tokens(self, tokens, incremental_state): if incremental_state is not None: # keep only the last token for incremental forward pass tokens = tokens[:, -1:] return self.embed_tokens(tokens) def _split_encoder_out(self, encoder_out, incremental_state): """Split and transpose encoder outputs. This is cached when doing incremental inference. """ cached_result = self.get_incremental_state('encoder_out') if cached_result: cached_result = utils.get_incremental_state(self, incremental_state, 'encoder_out') if cached_result is not None: return cached_result # transpose only once to speed up attention layers Loading @@ -320,12 +318,14 @@ class FConvDecoder(FairseqIncrementalDecoder): encoder_a = encoder_a.transpose(1, 2).contiguous() result = (encoder_a, encoder_b) return self.set_incremental_state('encoder_out', result) if incremental_state is not None: utils.set_incremental_state(self, incremental_state, 'encoder_out', result) return result def _transpose_unless_incremental_eval(self, x): if self._is_incremental_eval: def _transpose_if_training(self, x, incremental_state): if incremental_state is None: x = x.transpose(0, 1) return x return x.transpose(0, 1) def Embedding(num_embeddings, embedding_dim, padding_idx): Loading fairseq/models/lstm.py +24 −37 Changes for fairseq/models/lstm.py: 24 added lines, 37 removed lines. Original line number Diff line number Diff line Loading @@ -183,12 +183,9 @@ class LSTMDecoder(FairseqIncrementalDecoder): self.additional_fc = Linear(embed_dim, out_embed_dim) self.fc_out = Linear(out_embed_dim, num_embeddings, dropout=dropout_out) def forward(self, prev_output_tokens, encoder_out): if self._is_incremental_eval: def forward(self, prev_output_tokens, encoder_out, incremental_state=None): if incremental_state is not None: prev_output_tokens = prev_output_tokens[:, -1:] return self._forward(prev_output_tokens, encoder_out) def _forward(self, prev_output_tokens, encoder_out): bsz, seqlen = prev_output_tokens.size() # get outputs from encoder Loading @@ -204,15 +201,15 @@ class LSTMDecoder(FairseqIncrementalDecoder): x = x.transpose(0, 1) # initialize previous states (or get from cache during incremental generation) prev_hiddens = self.get_incremental_state('prev_hiddens') if not prev_hiddens: # first time step, initialize previous states prev_hiddens, prev_cells = self._init_prev_states(encoder_out) input_feed = Variable(x.data.new(bsz, embed_dim).zero_()) cached_state = utils.get_incremental_state(self, incremental_state, 'cached_state') if cached_state is not None: prev_hiddens, prev_cells, input_feed = cached_state else: # previous states are cached prev_cells = self.get_incremental_state('prev_cells') input_feed = self.get_incremental_state('input_feed') _, encoder_hiddens, encoder_cells = encoder_out num_layers = len(self.layers) prev_hiddens = [encoder_hiddens[i] for i in range(num_layers)] prev_cells = [encoder_cells[i] for i in range(num_layers)] input_feed = Variable(x.data.new(bsz, embed_dim).zero_()) attn_scores = Variable(x.data.new(srclen, seqlen, bsz).zero_()) outs = [] Loading Loading @@ -242,9 +239,8 @@ class LSTMDecoder(FairseqIncrementalDecoder): outs.append(out) # cache previous states (no-op except during incremental generation) self.set_incremental_state('prev_hiddens', prev_hiddens) self.set_incremental_state('prev_cells', prev_cells) self.set_incremental_state('input_feed', input_feed) utils.set_incremental_state( self, incremental_state, 'cached_state', (prev_hiddens, prev_cells, input_feed)) # collect outputs across time steps x = torch.cat(outs, dim=0).view(seqlen, bsz, embed_dim) Loading @@ -263,34 +259,25 @@ class LSTMDecoder(FairseqIncrementalDecoder): return x, attn_scores def reorder_incremental_state(self, new_order): """Reorder buffered internal state (for incremental generation).""" super().reorder_incremental_state(new_order) new_order = Variable(new_order) def reorder_incremental_state(self, incremental_state, new_order): cached_state = utils.get_incremental_state(self, incremental_state, 'cached_state') if cached_state is None: return def reorder_state(key): old = self.get_incremental_state(key) if isinstance(old, list): new = [old_i.index_select(0, new_order) for old_i in old] else: new = old.index_select(0, new_order) self.set_incremental_state(key, new) def reorder_state(state): if isinstance(state, list): return [reorder_state(state_i) for state_i in state] return state.index_select(0, new_order) reorder_state('prev_hiddens') reorder_state('prev_cells') reorder_state('input_feed') if not isinstance(new_order, Variable): new_order = Variable(new_order) new_state = tuple(map(reorder_state, cached_state)) utils.set_incremental_state(self, incremental_state, 'cached_state', new_state) def max_positions(self): """Maximum output length supported by the decoder.""" return int(1e5) # an arbitrary large number def _init_prev_states(self, encoder_out): _, encoder_hiddens, encoder_cells = encoder_out num_layers = len(self.layers) prev_hiddens = [encoder_hiddens[i] for i in range(num_layers)] prev_cells = [encoder_cells[i] for i in range(num_layers)] return prev_hiddens, prev_cells def Embedding(num_embeddings, embedding_dim, padding_idx): m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx) Loading fairseq/modules/learned_positional_embedding.py +2 −6 Changes for fairseq/modules/learned_positional_embedding.py: 2 added lines, 6 removed lines. Original line number Diff line number Diff line Loading @@ -20,14 +20,10 @@ class LearnedPositionalEmbedding(nn.Embedding): def __init__(self, num_embeddings, embedding_dim, padding_idx, left_pad): super().__init__(num_embeddings, embedding_dim, padding_idx) self.left_pad = left_pad self._is_incremental_eval = False def incremental_eval(self, mode=True): self._is_incremental_eval = mode def forward(self, input): def forward(self, input, incremental_state=None): """Input is expected to be of size [bsz x seqlen].""" if self._is_incremental_eval: if incremental_state is not None: # positions is the same for every token when decoding a single step positions = Variable( input.data.new(1, 1).fill_(self.padding_idx + input.size(1))) Loading fairseq/modules/linearized_convolution.py +32 −35 Changes for fairseq/modules/linearized_convolution.py: 32 added lines, 35 removed lines. Original line number Diff line number Diff line Loading @@ -22,35 +22,20 @@ class LinearizedConvolution(ConvTBC): def __init__(self, in_channels, out_channels, kernel_size, **kwargs): super().__init__(in_channels, out_channels, kernel_size, **kwargs) self._is_incremental_eval = False self._linearized_weight = None self.register_backward_hook(self._clear_linearized_weight) def remove_future_timesteps(self, x): """Remove future time steps created by padding.""" if not self._is_incremental_eval and self.kernel_size[0] > 1 and self.padding[0] > 0: x = x[:-self.padding[0], :, :] return x def incremental_eval(self, mode=True): self._is_incremental_eval = mode if mode: self.clear_incremental_state() def forward(self, input): if self._is_incremental_eval: return self.incremental_forward(input) else: def forward(self, input, incremental_state=None): """ Input: Time x Batch x Channel. Args: incremental_state: Used to buffer signal; if not None, then input is expected to contain a single frame. If the input order changes between time steps, call reorder_incremental_state. """ if incremental_state is None: return super().forward(input) def incremental_forward(self, input): """Forward convolution one time step at a time. This function maintains an internal state to buffer signal and accepts a single frame as input. If the input order changes between time steps, call reorder_incremental_state. To apply to fresh inputs, call clear_incremental_state. """ # reshape weight weight = self._get_linearized_weight() kw = self.kernel_size[0] Loading @@ -58,25 +43,37 @@ class LinearizedConvolution(ConvTBC): bsz = input.size(0) # input: bsz x len x dim if kw > 1: input = input.data if self.input_buffer is None: self.input_buffer = input.new(bsz, kw, input.size(2)) self.input_buffer.zero_() input_buffer = self._get_input_buffer(incremental_state) if input_buffer is None: input_buffer = input.new(bsz, kw, input.size(2)).zero_() self._set_input_buffer(incremental_state, input_buffer) else: # shift buffer self.input_buffer[:, :-1, :] = self.input_buffer[:, 1:, :].clone() input_buffer[:, :-1, :] = input_buffer[:, 1:, :].clone() # append next input self.input_buffer[:, -1, :] = input[:, -1, :] input = utils.volatile_variable(self.input_buffer) input_buffer[:, -1, :] = input[:, -1, :] input = utils.volatile_variable(input_buffer) with utils.maybe_no_grad(): output = F.linear(input.view(bsz, -1), weight, self.bias) return output.view(bsz, 1, -1) def clear_incremental_state(self): self.input_buffer = None def remove_future_timesteps(self, x): """Remove future time steps created by padding.""" if self.kernel_size[0] > 1 and self.padding[0] > 0: x = x[:-self.padding[0], :, :] return x def reorder_incremental_state(self, incremental_state, new_order): input_buffer = self._get_input_buffer(incremental_state) if input_buffer is not None: input_buffer = input_buffer.index_select(0, new_order) self._set_input_buffer(incremental_state, input_buffer) def _get_input_buffer(self, incremental_state): return utils.get_incremental_state(self, incremental_state, 'input_buffer') def reorder_incremental_state(self, new_order): if self.input_buffer is not None: self.input_buffer = self.input_buffer.index_select(0, new_order) def _set_input_buffer(self, incremental_state, new_buffer): return utils.set_incremental_state(self, incremental_state, 'input_buffer', new_buffer) def _get_linearized_weight(self): if self._linearized_weight is None: Loading Loading
fairseq/models/fairseq_incremental_decoder.py +5 −84 Changes for fairseq/models/fairseq_incremental_decoder.py: 5 added lines, 84 removed lines. Original line number Diff line number Diff line Loading @@ -13,99 +13,20 @@ class FairseqIncrementalDecoder(FairseqDecoder): def __init__(self, dictionary): super().__init__(dictionary) self._is_incremental_eval = False self._incremental_state = {} def forward(self, prev_output_tokens, encoder_out): if self._is_incremental_eval: def forward(self, prev_output_tokens, encoder_out, incremental_state=None): raise NotImplementedError else: raise NotImplementedError def incremental_inference(self): """Context manager for incremental inference. This provides an optimized forward pass for incremental inference (i.e., it predicts one time step at a time). If the input order changes between time steps, call reorder_incremental_state to update the relevant buffers. To generate a fresh sequence, first call clear_incremental_state. Usage: ``` with model.decoder.incremental_inference(): for step in range(maxlen): out, _ = model.decoder(tokens[:, :step], encoder_out) probs = model.get_normalized_probs(out[:, -1, :], log_probs=False) ``` """ class IncrementalInference(object): def __init__(self, decoder): self.decoder = decoder def __enter__(self): self.decoder.incremental_eval(True) def __exit__(self, *args): self.decoder.incremental_eval(False) return IncrementalInference(self) def incremental_eval(self, mode=True): """Sets the decoder and all children in incremental evaluation mode.""" assert self._is_incremental_eval != mode, \ 'incremental_eval already set to mode {}'.format(mode) self._is_incremental_eval = mode if mode: self.clear_incremental_state() def apply_incremental_eval(module): if module != self and hasattr(module, 'incremental_eval'): module.incremental_eval(mode) self.apply(apply_incremental_eval) def get_incremental_state(self, key): """Return cached state or None if not in incremental inference mode.""" if self._is_incremental_eval and key in self._incremental_state: return self._incremental_state[key] return None def set_incremental_state(self, key, value): """Cache state needed for incremental inference mode.""" if self._is_incremental_eval: self._incremental_state[key] = value return value def clear_incremental_state(self): """Clear all state used for incremental generation. **For incremental inference only** This should be called before generating a fresh sequence. beam_size is required if using BeamableMM. """ if self._is_incremental_eval: del self._incremental_state self._incremental_state = {} def apply_clear_incremental_state(module): if module != self and hasattr(module, 'clear_incremental_state'): module.clear_incremental_state() self.apply(apply_clear_incremental_state) def reorder_incremental_state(self, new_order): """Reorder buffered internal state (for incremental generation). **For incremental inference only** def reorder_incremental_state(self, incremental_state, new_order): """Reorder incremental state. This should be called when the order of the input has changed from the previous time step. A typical use case is beam search, where the input order changes between time steps based on the choice of beams. order changes between time steps based on the selection of beams. """ if self._is_incremental_eval: def apply_reorder_incremental_state(module): if module != self and hasattr(module, 'reorder_incremental_state'): module.reorder_incremental_state(new_order) module.reorder_incremental_state(incremental_state, new_order) self.apply(apply_reorder_incremental_state) def set_beam_size(self, beam_size): Loading
fairseq/models/fconv.py +27 −27 Changes for fairseq/models/fconv.py: 27 added lines, 27 removed lines. Original line number Diff line number Diff line Loading @@ -10,6 +10,7 @@ import torch import torch.nn as nn import torch.nn.functional as F from fairseq import utils from fairseq.data import LanguagePairDataset from fairseq.modules import BeamableMM, GradMultiply, LearnedPositionalEmbedding, LinearizedConvolution Loading Loading @@ -229,19 +230,13 @@ class FConvDecoder(FairseqIncrementalDecoder): else: self.fc3 = Linear(out_embed_dim, num_embeddings, dropout=dropout) def forward(self, prev_output_tokens, encoder_out): def forward(self, prev_output_tokens, encoder_out, incremental_state=None): # split and transpose encoder outputs encoder_a, encoder_b = self._split_encoder_out(encoder_out) encoder_a, encoder_b = self._split_encoder_out(encoder_out, incremental_state) # embed positions positions = self.embed_positions(prev_output_tokens) if self._is_incremental_eval: # keep only the last token for incremental forward pass prev_output_tokens = prev_output_tokens[:, -1:] # embed tokens and positions x = self.embed_tokens(prev_output_tokens) + positions # embed tokens and combine with positional embeddings x = self._embed_tokens(prev_output_tokens, incremental_state) x += self.embed_positions(prev_output_tokens, incremental_state) x = F.dropout(x, p=self.dropout, training=self.training) target_embedding = x Loading @@ -249,7 +244,7 @@ class FConvDecoder(FairseqIncrementalDecoder): x = self.fc1(x) # B x T x C -> T x B x C x = self._transpose_unless_incremental_eval(x) x = self._transpose_if_training(x, incremental_state) # temporal convolutions avg_attn_scores = None Loading @@ -258,13 +253,14 @@ class FConvDecoder(FairseqIncrementalDecoder): residual = x if proj is None else proj(x) x = F.dropout(x, p=self.dropout, training=self.training) x = conv(x) x = conv(x, incremental_state) if incremental_state is None: x = conv.remove_future_timesteps(x) x = F.glu(x, dim=2) # attention if attention is not None: x = self._transpose_unless_incremental_eval(x) x = self._transpose_if_training(x, incremental_state) x, attn_scores = attention(x, target_embedding, (encoder_a, encoder_b)) attn_scores = attn_scores / num_attn_layers Loading @@ -273,13 +269,13 @@ class FConvDecoder(FairseqIncrementalDecoder): else: avg_attn_scores.add_(attn_scores) x = self._transpose_unless_incremental_eval(x) x = self._transpose_if_training(x, incremental_state) # residual x = (x + residual) * math.sqrt(0.5) # T x B x C -> B x T x C x = self._transpose_unless_incremental_eval(x) x = self._transpose_if_training(x, incremental_state) # project back to size of vocabulary x = self.fc2(x) Loading @@ -288,10 +284,6 @@ class FConvDecoder(FairseqIncrementalDecoder): return x, avg_attn_scores def reorder_incremental_state(self, new_order): """Reorder buffered internal state (for incremental generation).""" super().reorder_incremental_state(new_order) def max_positions(self): """Maximum output length supported by the decoder.""" return self.embed_positions.max_positions() Loading @@ -306,13 +298,19 @@ class FConvDecoder(FairseqIncrementalDecoder): state_dict['decoder.version'] = torch.Tensor([1]) return state_dict def _split_encoder_out(self, encoder_out): def _embed_tokens(self, tokens, incremental_state): if incremental_state is not None: # keep only the last token for incremental forward pass tokens = tokens[:, -1:] return self.embed_tokens(tokens) def _split_encoder_out(self, encoder_out, incremental_state): """Split and transpose encoder outputs. This is cached when doing incremental inference. """ cached_result = self.get_incremental_state('encoder_out') if cached_result: cached_result = utils.get_incremental_state(self, incremental_state, 'encoder_out') if cached_result is not None: return cached_result # transpose only once to speed up attention layers Loading @@ -320,12 +318,14 @@ class FConvDecoder(FairseqIncrementalDecoder): encoder_a = encoder_a.transpose(1, 2).contiguous() result = (encoder_a, encoder_b) return self.set_incremental_state('encoder_out', result) if incremental_state is not None: utils.set_incremental_state(self, incremental_state, 'encoder_out', result) return result def _transpose_unless_incremental_eval(self, x): if self._is_incremental_eval: def _transpose_if_training(self, x, incremental_state): if incremental_state is None: x = x.transpose(0, 1) return x return x.transpose(0, 1) def Embedding(num_embeddings, embedding_dim, padding_idx): Loading
fairseq/models/lstm.py +24 −37 Changes for fairseq/models/lstm.py: 24 added lines, 37 removed lines. Original line number Diff line number Diff line Loading @@ -183,12 +183,9 @@ class LSTMDecoder(FairseqIncrementalDecoder): self.additional_fc = Linear(embed_dim, out_embed_dim) self.fc_out = Linear(out_embed_dim, num_embeddings, dropout=dropout_out) def forward(self, prev_output_tokens, encoder_out): if self._is_incremental_eval: def forward(self, prev_output_tokens, encoder_out, incremental_state=None): if incremental_state is not None: prev_output_tokens = prev_output_tokens[:, -1:] return self._forward(prev_output_tokens, encoder_out) def _forward(self, prev_output_tokens, encoder_out): bsz, seqlen = prev_output_tokens.size() # get outputs from encoder Loading @@ -204,15 +201,15 @@ class LSTMDecoder(FairseqIncrementalDecoder): x = x.transpose(0, 1) # initialize previous states (or get from cache during incremental generation) prev_hiddens = self.get_incremental_state('prev_hiddens') if not prev_hiddens: # first time step, initialize previous states prev_hiddens, prev_cells = self._init_prev_states(encoder_out) input_feed = Variable(x.data.new(bsz, embed_dim).zero_()) cached_state = utils.get_incremental_state(self, incremental_state, 'cached_state') if cached_state is not None: prev_hiddens, prev_cells, input_feed = cached_state else: # previous states are cached prev_cells = self.get_incremental_state('prev_cells') input_feed = self.get_incremental_state('input_feed') _, encoder_hiddens, encoder_cells = encoder_out num_layers = len(self.layers) prev_hiddens = [encoder_hiddens[i] for i in range(num_layers)] prev_cells = [encoder_cells[i] for i in range(num_layers)] input_feed = Variable(x.data.new(bsz, embed_dim).zero_()) attn_scores = Variable(x.data.new(srclen, seqlen, bsz).zero_()) outs = [] Loading Loading @@ -242,9 +239,8 @@ class LSTMDecoder(FairseqIncrementalDecoder): outs.append(out) # cache previous states (no-op except during incremental generation) self.set_incremental_state('prev_hiddens', prev_hiddens) self.set_incremental_state('prev_cells', prev_cells) self.set_incremental_state('input_feed', input_feed) utils.set_incremental_state( self, incremental_state, 'cached_state', (prev_hiddens, prev_cells, input_feed)) # collect outputs across time steps x = torch.cat(outs, dim=0).view(seqlen, bsz, embed_dim) Loading @@ -263,34 +259,25 @@ class LSTMDecoder(FairseqIncrementalDecoder): return x, attn_scores def reorder_incremental_state(self, new_order): """Reorder buffered internal state (for incremental generation).""" super().reorder_incremental_state(new_order) new_order = Variable(new_order) def reorder_incremental_state(self, incremental_state, new_order): cached_state = utils.get_incremental_state(self, incremental_state, 'cached_state') if cached_state is None: return def reorder_state(key): old = self.get_incremental_state(key) if isinstance(old, list): new = [old_i.index_select(0, new_order) for old_i in old] else: new = old.index_select(0, new_order) self.set_incremental_state(key, new) def reorder_state(state): if isinstance(state, list): return [reorder_state(state_i) for state_i in state] return state.index_select(0, new_order) reorder_state('prev_hiddens') reorder_state('prev_cells') reorder_state('input_feed') if not isinstance(new_order, Variable): new_order = Variable(new_order) new_state = tuple(map(reorder_state, cached_state)) utils.set_incremental_state(self, incremental_state, 'cached_state', new_state) def max_positions(self): """Maximum output length supported by the decoder.""" return int(1e5) # an arbitrary large number def _init_prev_states(self, encoder_out): _, encoder_hiddens, encoder_cells = encoder_out num_layers = len(self.layers) prev_hiddens = [encoder_hiddens[i] for i in range(num_layers)] prev_cells = [encoder_cells[i] for i in range(num_layers)] return prev_hiddens, prev_cells def Embedding(num_embeddings, embedding_dim, padding_idx): m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx) Loading
fairseq/modules/learned_positional_embedding.py +2 −6 Changes for fairseq/modules/learned_positional_embedding.py: 2 added lines, 6 removed lines. Original line number Diff line number Diff line Loading @@ -20,14 +20,10 @@ class LearnedPositionalEmbedding(nn.Embedding): def __init__(self, num_embeddings, embedding_dim, padding_idx, left_pad): super().__init__(num_embeddings, embedding_dim, padding_idx) self.left_pad = left_pad self._is_incremental_eval = False def incremental_eval(self, mode=True): self._is_incremental_eval = mode def forward(self, input): def forward(self, input, incremental_state=None): """Input is expected to be of size [bsz x seqlen].""" if self._is_incremental_eval: if incremental_state is not None: # positions is the same for every token when decoding a single step positions = Variable( input.data.new(1, 1).fill_(self.padding_idx + input.size(1))) Loading
fairseq/modules/linearized_convolution.py +32 −35 Changes for fairseq/modules/linearized_convolution.py: 32 added lines, 35 removed lines. Original line number Diff line number Diff line Loading @@ -22,35 +22,20 @@ class LinearizedConvolution(ConvTBC): def __init__(self, in_channels, out_channels, kernel_size, **kwargs): super().__init__(in_channels, out_channels, kernel_size, **kwargs) self._is_incremental_eval = False self._linearized_weight = None self.register_backward_hook(self._clear_linearized_weight) def remove_future_timesteps(self, x): """Remove future time steps created by padding.""" if not self._is_incremental_eval and self.kernel_size[0] > 1 and self.padding[0] > 0: x = x[:-self.padding[0], :, :] return x def incremental_eval(self, mode=True): self._is_incremental_eval = mode if mode: self.clear_incremental_state() def forward(self, input): if self._is_incremental_eval: return self.incremental_forward(input) else: def forward(self, input, incremental_state=None): """ Input: Time x Batch x Channel. Args: incremental_state: Used to buffer signal; if not None, then input is expected to contain a single frame. If the input order changes between time steps, call reorder_incremental_state. """ if incremental_state is None: return super().forward(input) def incremental_forward(self, input): """Forward convolution one time step at a time. This function maintains an internal state to buffer signal and accepts a single frame as input. If the input order changes between time steps, call reorder_incremental_state. To apply to fresh inputs, call clear_incremental_state. """ # reshape weight weight = self._get_linearized_weight() kw = self.kernel_size[0] Loading @@ -58,25 +43,37 @@ class LinearizedConvolution(ConvTBC): bsz = input.size(0) # input: bsz x len x dim if kw > 1: input = input.data if self.input_buffer is None: self.input_buffer = input.new(bsz, kw, input.size(2)) self.input_buffer.zero_() input_buffer = self._get_input_buffer(incremental_state) if input_buffer is None: input_buffer = input.new(bsz, kw, input.size(2)).zero_() self._set_input_buffer(incremental_state, input_buffer) else: # shift buffer self.input_buffer[:, :-1, :] = self.input_buffer[:, 1:, :].clone() input_buffer[:, :-1, :] = input_buffer[:, 1:, :].clone() # append next input self.input_buffer[:, -1, :] = input[:, -1, :] input = utils.volatile_variable(self.input_buffer) input_buffer[:, -1, :] = input[:, -1, :] input = utils.volatile_variable(input_buffer) with utils.maybe_no_grad(): output = F.linear(input.view(bsz, -1), weight, self.bias) return output.view(bsz, 1, -1) def clear_incremental_state(self): self.input_buffer = None def remove_future_timesteps(self, x): """Remove future time steps created by padding.""" if self.kernel_size[0] > 1 and self.padding[0] > 0: x = x[:-self.padding[0], :, :] return x def reorder_incremental_state(self, incremental_state, new_order): input_buffer = self._get_input_buffer(incremental_state) if input_buffer is not None: input_buffer = input_buffer.index_select(0, new_order) self._set_input_buffer(incremental_state, input_buffer) def _get_input_buffer(self, incremental_state): return utils.get_incremental_state(self, incremental_state, 'input_buffer') def reorder_incremental_state(self, new_order): if self.input_buffer is not None: self.input_buffer = self.input_buffer.index_select(0, new_order) def _set_input_buffer(self, incremental_state, new_buffer): return utils.set_incremental_state(self, incremental_state, 'input_buffer', new_buffer) def _get_linearized_weight(self): if self._linearized_weight is None: Loading
