add arg num_workers for ernie3.0

PaddlePaddle · Apr 10, 2023 · dc08470 · dc08470
1 parent c2f7e31
commit dc08470
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/tests/test_tipc/benchmark/modules/ernie3_for_sequence_classification.py b/tests/test_tipc/benchmark/modules/ernie3_for_sequence_classification.py
@@ -145,7 +145,7 @@ def create_data_loader(self, args, **kwargs):
         train_loader = DataLoader(
             dataset=train_ds,
             batch_sampler=train_batch_sampler,
-            num_workers=4,  # when paddlepaddle<=2.4.1, if we use dynamicTostatic mode, we need set num_workeks > 0
+            num_workers=args.num_workers,  # when paddlepaddle<=2.4.1, if we use dynamicTostatic mode, we need set num_workeks > 0
         )
 
         self.num_batch = len(train_loader)

diff --git a/tests/test_tipc/benchmark/options.py b/tests/test_tipc/benchmark/options.py
@@ -136,6 +136,12 @@ def get_parser():
     parser.add_argument("--epoch", type=int, default=10, help="Number of epochs. ")
 
     parser.add_argument("--generated_inputs", action="store_true", help="Use generated inputs. ")
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=4,
+        help="num_workers of dataloader. When paddlepaddle<=2.4.1, if we use dynamicTostatic mode, we need set num_workeks > 0 ",
+    )
 
     # For benchmark.
     parser.add_argument(