Make FT selectable (PaddlePaddle#826)

* Make FT selectable

* update

* fix comments

Author: FrostML

examples/machine_translation/transformer/predict.py

@@ -30,6 +30,10 @@ def parse_args():
         type=str,
         help="The file for testing. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used to process testing."
     )
+    parser.add_argument(
+        "--without_ft",
+        action="store_true",
+        help="Whether to use Faster Transformer to do predict. ")
     args = parser.parse_args()
     return args
 
@@ -78,7 +82,8 @@ def do_predict(args):
         bos_id=args.bos_idx,
         eos_id=args.eos_idx,
         beam_size=args.beam_size,
-        max_out_len=args.max_out_len)
+        max_out_len=args.max_out_len,
+        use_ft=not args.without_ft)
 
     # Load the trained model
     assert args.init_from_params, (
@@ -114,6 +119,7 @@ def do_predict(args):
         args = AttrDict(yaml.safe_load(f))
     args.benchmark = ARGS.benchmark
     args.test_file = ARGS.test_file
+    args.without_ft = ARGS.without_ft
     pprint(args)
 
     do_predict(args)

paddlenlp/ops/faster_transformer/transformer/faster_transformer.py

@@ -268,13 +268,14 @@ class TransformerGenerator(paddle.nn.Layer):
         max_out_len (int, optional):
             The maximum output length. Defaults to 256.
         kwargs:
-            The key word arguments can be `output_time_major` and `use_fp16_decoding`.
+            The key word arguments can be `output_time_major`, `use_fp16_decoding` and `use_ft`.
             `output_time_major(bool, optional)`: Indicate the data layout of predicted
             Tensor. If `False`, the data layout would be batch major with shape
             `[batch_size, seq_len, beam_size]`. If  `True`, the data layout would
             be time major with shape `[seq_len, batch_size, beam_size]`. Default
             to `False`. `use_fp16_decoding(bool, optional)`: Whether to use fp16
-            for decoding.
+            for decoding. `use_ft(bool, optional)`: Whether to use Faster Transformer
+            for decoding. 
     """
 
     def __init__(self,
@@ -303,25 +304,48 @@ def __init__(self,
         self.max_length = max_length
         self.output_time_major = kwargs.pop("output_time_major", True)
         use_fp16_decoding = kwargs.pop("use_fp16_decoding", False)
-        try:
-            load("FasterTransformer", verbose=True)
-            self.transformer = FasterTransformer(
-                src_vocab_size=src_vocab_size,
-                trg_vocab_size=trg_vocab_size,
-                max_length=max_length,
-                num_encoder_layers=num_encoder_layers,
-                num_decoder_layers=num_decoder_layers,
-                n_head=n_head,
-                d_model=d_model,
-                d_inner_hid=d_inner_hid,
-                dropout=dropout,
-                weight_sharing=weight_sharing,
-                bos_id=bos_id,
-                eos_id=eos_id,
-                beam_size=beam_size,
-                max_out_len=max_out_len,
-                use_fp16_decoding=use_fp16_decoding)
-        except Exception:
+        use_ft = kwargs.pop("use_ft", True)
+
+        if use_ft:
+            try:
+                load("FasterTransformer", verbose=True)
+                self.transformer = FasterTransformer(
+                    src_vocab_size=src_vocab_size,
+                    trg_vocab_size=trg_vocab_size,
+                    max_length=max_length,
+                    num_encoder_layers=num_encoder_layers,
+                    num_decoder_layers=num_decoder_layers,
+                    n_head=n_head,
+                    d_model=d_model,
+                    d_inner_hid=d_inner_hid,
+                    dropout=dropout,
+                    weight_sharing=weight_sharing,
+                    bos_id=bos_id,
+                    eos_id=eos_id,
+                    beam_size=beam_size,
+                    max_out_len=max_out_len,
+                    use_fp16_decoding=use_fp16_decoding)
+            except Exception:
+                logger.warning(
+                    "Exception occurs when using Faster Transformer. " \
+                    "The original forward will be involved. ")
+                self.transformer = InferTransformerModel(
+                    src_vocab_size=src_vocab_size,
+                    trg_vocab_size=trg_vocab_size,
+                    max_length=max_length,
+                    num_encoder_layers=num_encoder_layers,
+                    num_decoder_layers=num_decoder_layers,
+                    n_head=n_head,
+                    d_model=d_model,
+                    d_inner_hid=d_inner_hid,
+                    dropout=dropout,
+                    weight_sharing=weight_sharing,
+                    bos_id=bos_id,
+                    eos_id=eos_id,
+                    beam_size=beam_size,
+                    max_out_len=max_out_len,
+                    output_time_major=self.output_time_major)
+        else:
             self.transformer = InferTransformerModel(
                 src_vocab_size=src_vocab_size,
                 trg_vocab_size=trg_vocab_size,