Implement repeat and beam_update via map_batch_fn

onmt/decoders/cnn_decoder.py

@@ -137,20 +137,14 @@ def __init__(self, memory_bank, enc_hidden):
         self.init_src = (memory_bank + enc_hidden) * SCALE_WEIGHT
         self.previous_input = None
 
-    @property
-    def _all(self):
-        """
-        Contains attributes that need to be updated in self.beam_update().
-        """
-        return (self.previous_input,)
-
     def detach(self):
         self.previous_input = self.previous_input.detach()
 
     def update_state(self, new_input):
         """ Called for every decoder forward pass. """
         self.previous_input = new_input
 
-    def repeat_beam_size_times(self, beam_size):
-        """ Repeat beam_size times along batch dimension. """
-        self.init_src = self.init_src.data.repeat(1, beam_size, 1)
+    def map_batch_fn(self, fn):
+        self.init_src = fn(self.init_src, 1)
+        if self.previous_input is not None:
+            self.previous_input = fn(self.previous_input, 1)

onmt/decoders/decoder.py

@@ -390,23 +390,6 @@ def detach(self):
         self.hidden = tuple([_.detach() for _ in self.hidden])
         self.input_feed = self.input_feed.detach()
 
-    def beam_update(self, idx, positions, beam_size):
-        """ Need to document this """
-        for e in self._all:
-            sizes = e.size()
-            br = sizes[1]
-            if len(sizes) == 3:
-                sent_states = e.view(sizes[0], beam_size, br // beam_size,
-                                     sizes[2])[:, :, idx]
-            else:
-                sent_states = e.view(sizes[0], beam_size,
-                                     br // beam_size,
-                                     sizes[2],
-                                     sizes[3])[:, :, idx]
-
-            sent_states.data.copy_(
-                sent_states.data.index_select(1, positions))
-
     def map_batch_fn(self, fn):
         raise NotImplementedError()
 
@@ -433,10 +416,6 @@ def __init__(self, hidden_size, rnnstate):
         self.input_feed = self.hidden[0].data.new(*h_size).zero_() \
                               .unsqueeze(0)
 
-    @property
-    def _all(self):
-        return self.hidden + (self.input_feed,)
-
     def update_state(self, rnnstate, input_feed, coverage):
         """ Update decoder state """
         if not isinstance(rnnstate, tuple):
@@ -446,13 +425,6 @@ def update_state(self, rnnstate, input_feed, coverage):
         self.input_feed = input_feed
         self.coverage = coverage
 
-    def repeat_beam_size_times(self, beam_size):
-        """ Repeat beam_size times along batch dimension. """
-        vars = [e.data.repeat(1, beam_size, 1)
-                for e in self._all]
-        self.hidden = tuple(vars[:-1])
-        self.input_feed = vars[-1]
-
     def map_batch_fn(self, fn):
         self.hidden = tuple(map(lambda x: fn(x, 1), self.hidden))
         self.input_feed = fn(self.input_feed, 1)

onmt/decoders/ensemble.py

@@ -20,14 +20,9 @@ class EnsembleDecoderState(DecoderState):
     def __init__(self, model_decoder_states):
         self.model_decoder_states = tuple(model_decoder_states)
 
-    def beam_update(self, idx, positions, beam_size):
+    def map_batch_fn(self, fn):
         for model_state in self.model_decoder_states:
-            model_state.beam_update(idx, positions, beam_size)
-
-    def repeat_beam_size_times(self, beam_size):
-        """ Repeat beam_size times along batch dimension. """
-        for model_state in self.model_decoder_states:
-            model_state.repeat_beam_size_times(beam_size)
+            model_state.map_batch_fn(fn)
 
     def __getitem__(self, index):
         return self.model_decoder_states[index]

onmt/decoders/transformer.py

@@ -260,19 +260,6 @@ def __init__(self, src):
         self.previous_layer_inputs = None
         self.cache = None
 
-    @property
-    def _all(self):
-        """
-        Contains attributes that need to be updated in self.beam_update().
-        """
-        if (self.previous_input is not None
-                and self.previous_layer_inputs is not None):
-            return (self.previous_input,
-                    self.previous_layer_inputs,
-                    self.src)
-        else:
-            return (self.src,)
-
     def detach(self):
         if self.previous_input is not None:
             self.previous_input = self.previous_input.detach()
@@ -306,10 +293,6 @@ def _init_cache(self, memory_bank, num_layers, self_attn_type):
                 layer_cache["self_values"] = None
             self.cache["layer_{}".format(l)] = layer_cache
 
-    def repeat_beam_size_times(self, beam_size):
-        """ Repeat beam_size times along batch dimension. """
-        self.src = self.src.data.repeat(1, beam_size, 1)
-
     def map_batch_fn(self, fn):
         def _recursive_map(struct, batch_dim=0):
             for k, v in struct.items():
@@ -320,5 +303,9 @@ def _recursive_map(struct, batch_dim=0):
                         struct[k] = fn(v, batch_dim)
 
         self.src = fn(self.src, 1)
+        if self.previous_input is not None:
+            self.previous_input = fn(self.previous_input, 1)
+        if self.previous_layer_inputs is not None:
+            self.previous_layer_inputs = fn(self.previous_layer_inputs, 1)
         if self.cache is not None:
             _recursive_map(self.cache)

onmt/translate/translator.py

@@ -566,6 +566,11 @@ def bottle(m):
         def unbottle(m):
             return m.view(beam_size, batch_size, -1)
 
+        def _repeat_beam_size_times(x, dim):
+            repeats = [1] * x.dim()
+            repeats[dim] = beam_size
+            return x.repeat(*repeats)
+
         # (1) Run the encoder on the src.
         src, enc_states, memory_bank, src_lengths = self._run_encoder(
             batch, data_type)
@@ -580,7 +585,7 @@ def unbottle(m):
         else:
             memory_bank = rvar(memory_bank.data)
         memory_lengths = src_lengths.repeat(beam_size)
-        dec_states.repeat_beam_size_times(beam_size)
+        dec_states.map_batch_fn(_repeat_beam_size_times)
 
         # (3) run the decoder to generate sentences, using beam search.
         for i in range(self.max_length):
@@ -631,10 +636,19 @@ def unbottle(m):
                 beam_attn = unbottle(attn["copy"])
 
             # (c) Advance each beam.
+            select_indices_array = []
             for j, b in enumerate(beam):
                 b.advance(out[:, j],
                           beam_attn.data[:, j, :memory_lengths[j]])
-                dec_states.beam_update(j, b.get_current_origin(), beam_size)
+                select_indices_array.append(
+                    b.get_current_origin() * batch_size + j)
+            select_indices = torch.cat(select_indices_array) \
+                                  .view(batch_size, beam_size) \
+                                  .transpose(0, 1) \
+                                  .contiguous() \
+                                  .view(-1)
+            dec_states.map_batch_fn(
+                lambda state, dim: state.index_select(dim, select_indices))
 
         # (4) Extract sentences from beam.
         ret = self._from_beam(beam)