def reorder_incremental_state(self, incremental_state, new_order): super().reorder_incremental_state(incremental_state, new_order) cached_state = utils.get_incremental_state(self, incremental_state, 'cached_state') if cached_state is None: return def reorder_state(state): if isinstance(state, list): return [reorder_state(state_i) for state_i in state] elif state is not None: return state.index_select(0, new_order) else: return None new_state = tuple(map(reorder_state, cached_state)) utils.set_incremental_state(self, incremental_state, 'cached_state', new_state)
def extract_features( self, prev_output_tokens, encoder_out, incremental_state=None ): """ Similar to *forward* but only return features. """ encoder_padding_masks, encoder_outs = {}, {} for modality in self.src_modalities: if encoder_out[modality] is not None: encoder_padding_masks[modality] = encoder_out[modality]['encoder_padding_mask'] encoder_outs[modality] = encoder_out[modality]['encoder_out'] else: encoder_padding_masks[modality] = None encoder_outs[modality] = None if incremental_state is not None: prev_output_tokens = prev_output_tokens[:, -1:] bsz, seqlen = prev_output_tokens.size() # get outputs from encoder encoder_outputs, encoder_hiddens, encoder_cells, srclen = {}, {}, {}, {} for modality in self.src_modalities: if encoder_outs[modality] is not None: encoder_outputs[modality], encoder_hiddens[modality], encoder_cells[modality] = encoder_outs[modality][ :3] srclen[modality] = encoder_outputs[modality].size(0) flag = True else: srclen[modality] = None if flag: # TODO # concatenate encoder_hiddens = torch.cat([encoder_hiddens[modality]] for modality in self.src_modalities) encoder_cells = torch.cat([encoder_cells[modality]] for modality in self.src_modalities) # embed tokens x = self.embed_tokens(prev_output_tokens) x = F.dropout(x, p=self.dropout_in, training=self.training) # B x T x C -> T x B x C x = x.transpose(0, 1) # initialize previous states (or get from cache during incremental generation) 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 elif encoder_out is not None: # setup recurrent cells 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)] if self.encoder_hidden_proj is not None: prev_hiddens = [self.encoder_hidden_proj(x) for x in prev_hiddens] prev_cells = [self.encoder_cell_proj(x) for x in prev_cells] input_feed = x.new_zeros(bsz, self.hidden_size) else: # setup zero cells, since there is no encoder num_layers = len(self.layers) zero_state = x.new_zeros(bsz, self.hidden_size) prev_hiddens = [zero_state for i in range(num_layers)] prev_cells = [zero_state for i in range(num_layers)] input_feed = None assert srclen is not None or self.attention is None, \ "attention is not supported if there are no encoder outputs" attn_scores = x.new_zeros(srclen, seqlen, bsz) if self.attention is not None else None outs = [] for j in range(seqlen): # input feeding: concatenate context vector from previous time step if input_feed is not None: input = torch.cat((x[j, :, :], input_feed), dim=1) else: input = x[j] for i, rnn in enumerate(self.layers): # recurrent cell hidden, cell = rnn(input, (prev_hiddens[i], prev_cells[i])) # hidden state becomes the input to the next layer input = F.dropout(hidden, p=self.dropout_out, training=self.training) # save state for next time step prev_hiddens[i] = hidden prev_cells[i] = cell # apply attention using the last layer's hidden state if self.attention is not None: pass # out, attn_scores[:, j, :] = self.attention(hidden, encoder_outs, encoder_padding_mask) else: out = hidden out = F.dropout(out, p=self.dropout_out, training=self.training) # input feeding if input_feed is not None: input_feed = out # save final output outs.append(out) # cache previous states (no-op except during incremental generation) 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, self.hidden_size) # T x B x C -> B x T x C x = x.transpose(1, 0) if hasattr(self, 'additional_fc') and self.adaptive_softmax is None: x = self.additional_fc(x) x = F.dropout(x, p=self.dropout_out, training=self.training) # srclen x tgtlen x bsz -> bsz x tgtlen x srclen if not self.training and self.need_attn and self.attention is not None: attn_scores = attn_scores.transpose(0, 2) else: attn_scores = None return x, attn_scores
def generate(self, models, sample, **kwargs): """Generate a batch of translations. Args: models (List[~fairseq.models.NccModel]): ensemble of models sample (dict): batch prefix_tokens (torch.LongTensor, optional): force decoder to begin with these tokens bos_token (int, optional): beginning of sentence token (default: self.eos) """ model = models[0] # for ensemble expansion if not self.retain_dropout: model.eval() # model.forward normally channels prev_output_tokens into the decoder # separately, but SequenceGenerator directly calls model.encoder encoder_input = { k: v for k, v in sample['net_input'].items() if k != 'prev_output_tokens' } src_tokens = encoder_input['src_tokens'] # batch dimension goes first followed by source lengths bsz = src_tokens[0].size(0) max_len = self.max_len_b assert self.min_len <= max_len, 'min_len cannot be larger than max_len, please adjust these!' # 1. encoder encoder_out = model.encoder( sample['net_input']['src_tokens'], src_lengths=sample['net_input']['src_lengths'], **kwargs) encoder_padding_mask = encoder_out['encoder_padding_mask'] encoder_out = encoder_out['encoder_out'] device = encoder_out[0].device prev_output_tokens = torch.zeros(bsz, 1).long().fill_(self.eos).to( device) # <eos> incremental_state = None # get outputs from encoder if encoder_out is not None: encoder_outs, encoder_hiddens, encoder_cells = encoder_out[:3] srclen = encoder_outs.size(1) encoder_outs = encoder_outs.transpose(0, 1) else: srclen = None # initialize previous states (or get from cache during incremental generation) cached_state = utils.get_incremental_state(model.decoder, incremental_state, 'cached_state') if cached_state is not None: prev_hiddens, prev_cells, input_feed = cached_state elif encoder_out is not None: # setup recurrent cells prev_hiddens = [encoder_hiddens] prev_cells = [encoder_cells] # equals to torch.zeros(bsz, model.decoder.hidden_size).to(device).type_as(encoder_out[0]), # but for support float16, we recommend such implementation # input_feed = encoder_out[0].new(bsz, model.decoder.hidden_size).fill_(0) input_feed = None else: # setup zero cells, since there is no encoder num_layers = len(model.decoder.layers) # for support float16 # zero_state = torch.zeros(bsz, model.decoder.hidden_size).to(device).type_as(encoder_out[0]) zero_state = encoder_out[0].new(bsz, model.decoder.hidden_size).fill_(0) prev_hiddens = [zero_state for i in range(num_layers)] prev_cells = [zero_state for i in range(num_layers)] input_feed = None assert srclen is not None or model.decoder.attention is None, \ "attention is not supported if there are no encoder outputs" # attn_scores = torch.zeros(srclen, max_len, bsz).to(device) if model.decoder.attention is not None else None attn_scores = encoder_out[0].new( srclen, max_len, bsz).fill_(0) if model.decoder.attention is not None else None dec_preds = [] # 2. generate for j in range(max_len): # embed tokens prev_output_tokens_emb = model.decoder.embed_tokens( prev_output_tokens) # B x T x C -> T x B x C prev_output_tokens_emb = prev_output_tokens_emb.squeeze( 1) # transpose(0, 1) # input feeding: concatenate context vector from previous time step if input_feed is not None: input = torch.cat((prev_output_tokens_emb, input_feed), dim=1) else: input = prev_output_tokens_emb for i, rnn in enumerate(model.decoder.layers): # recurrent cell hidden, cell = rnn(input, (prev_hiddens[i], prev_cells[i])) # save state for next time step prev_hiddens[i] = hidden prev_cells[i] = cell # apply attention using the last layer's hidden state if model.decoder.attention is not None: out, attn_scores[:, j, :] = model.decoder.attention( hidden, encoder_outs, encoder_padding_mask) else: out = hidden # input feeding if input_feed is not None: input_feed = out decoded = model.decoder.output_layer( out) # (batch_size*comment_dict_size) logprobs = F.log_softmax(decoded, dim=-1) # (batch_size*comment_dict_size) prob_prev = torch.exp(logprobs) # (batch_size*comment_dict_size) # input feeding if input_feed is not None: input_feed = out sample_max = True if sample_max: sample_logprobs, predicted = torch.max(prob_prev, 1) dec_preds.append(predicted.clone()) else: predicted = torch.multinomial(prob_prev, 1) # .to(device) dec_preds.append(predicted.clone()) # cache previous states (no-op except during incremental generation) utils.set_incremental_state( self, incremental_state, 'cached_state', (prev_hiddens, prev_cells, input_feed), ) prev_output_tokens = predicted.reshape(-1, 1) dec_preds = torch.stack(dec_preds, dim=1) predictions = [] for pred in dec_preds.tolist(): predictions.append([{ 'tokens': torch.Tensor(pred).type_as(dec_preds) }]) return predictions