* support OP fusion based on probed speed of each OP

modelscope · Sep 18, 2024 · f6e017f · f6e017f
1 parent 5a6c071
commit f6e017f
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/data_juicer/core/adapter.py b/data_juicer/core/adapter.py
@@ -31,6 +31,11 @@ def execute_and_probe(dataset, operators, sample_interval=0.5):
         resource_util_list = []
         # probe for each OP
         for op in operators:
+            # set num_proc to 1 for each OP to focus on the influence of batch
+            # size only.
+            old_num_proc = op.num_proc
+            op.num_proc = 1
+
             # number of test samples
             sample_num = len(dataset)
             # run single op and monitor the resource utilization
@@ -42,6 +47,9 @@ def execute_and_probe(dataset, operators, sample_interval=0.5):
                 'speed'] = sample_num / resource_util_per_op['time']
             resource_util_list.append(resource_util_per_op)
 
+            # restore to the original num_proc
+            op.num_proc = old_num_proc
+
         return resource_util_list
 
     @staticmethod