turned on bucketing

sockeye/score.py

@@ -61,7 +61,8 @@ def check_arg_compatibility(args: argparse.Namespace):
 def create_scoring_model(config: model.ModelConfig,
                          model_dir: str,
                          context: List[mx.Context],
-                         score_iter: data_io.BaseParallelSampleIter) -> scoring.ScoringModel:
+                         score_iter: data_io.BaseParallelSampleIter,
+                         bucketing: bool = False) -> scoring.ScoringModel:
     """
     Create a scoring model and load the parameters from disk if needed.
 
@@ -76,7 +77,8 @@ def create_scoring_model(config: model.ModelConfig,
                                          model_dir=model_dir,
                                          context=context,
                                          provide_data=score_iter.provide_data,
-                                         default_bucket_key=score_iter.default_bucket_key)
+                                         default_bucket_key=score_iter.default_bucket_key,
+                                         bucketing=bucketing)
 
     return scoring_model
 
@@ -132,6 +134,7 @@ def score(args: argparse.Namespace):
         scoring_model = create_scoring_model(config=model_config,
                                              model_dir=args.model,
                                              context=context,
+                                             bucketing=not args.no_bucketing,
                                              score_iter=score_iter)
 
         scorer = scoring.Scorer(scoring_model, source_vocabs, target_vocab,

sockeye/scoring.py

@@ -62,9 +62,11 @@ def __init__(self,
                  model_dir: str,
                  context: List[mx.context.Context],
                  provide_data: List[mx.io.DataDesc],
+                 bucketing: bool,
                  default_bucket_key: Tuple[int, int]) -> None:
         super().__init__(config)
         self.context = context
+        self.bucketing = bucketing
         self._initialize(provide_data, default_bucket_key)
 
         params_fname = os.path.join(model_dir, C.PARAMS_BEST_NAME)
@@ -134,16 +136,24 @@ def sym_gen(seq_lens):
             # return the outputs and the data names (we don't need the labels)
             return outputs, data_names, None
 
-        logger.info("Using bucketing. Default max_seq_len=%s", default_bucket_key)
-        self.module = mx.mod.BucketingModule(sym_gen=sym_gen,
-                                             logger=logger,
-                                             default_bucket_key=default_bucket_key,
-                                             context=self.context)
+        if self.bucketing:
+            logger.info("Using bucketing. Default max_seq_len=%s", default_bucket_key)
+            self.module = mx.mod.BucketingModule(sym_gen=sym_gen,
+                                                 logger=logger,
+                                                 default_bucket_key=default_bucket_key,
+                                                 context=self.context)
+        else:
+            symbol, _, __ = sym_gen(default_bucket_key)
+            self.module = mx.mod.Module(symbol=symbol,
+                                        data_names=data_names,
+                                        label_names=None,
+                                        logger=logger,
+                                        context=self.context)
 
         self.module.bind(data_shapes=provide_data,
                          label_shapes=None,
                          for_training=False,
-                         force_rebind=True,
+                         force_rebind=False,
                          grad_req=None)
 
 
@@ -207,9 +217,10 @@ def score(self,
               score_iter):
 
         for i, batch in enumerate(score_iter):
-            # data_io generates labels, too, which we don't need
-            label, batch.provide_label = batch.provide_label, None
+            # data_io generates labels, too, which aren't needed in the computation graph
+            batch.provide_label = None
             labels = batch.label[0].as_in_context(self.model.context[0])
+            batch.label = None
             self.model.prepare_batch(batch)
             self.model.run_forward(batch)
             outputs = self.model.get_outputs()