|
@@ -1,5 +1,6 @@
|
|
|
import argparse
|
|
import argparse
|
|
|
import sys
|
|
import sys
|
|
|
|
|
+import os
|
|
|
import time
|
|
import time
|
|
|
import traceback
|
|
import traceback
|
|
|
from multiprocessing import Manager, Process
|
|
from multiprocessing import Manager, Process
|
|
@@ -11,6 +12,15 @@ from paddlex.utils.device import constr_device, parse_device
|
|
|
|
|
|
|
|
|
|
|
|
|
def worker(pipeline_name_or_config_path, device, task_queue, batch_size, output_dir):
|
|
def worker(pipeline_name_or_config_path, device, task_queue, batch_size, output_dir):
|
|
|
|
|
+ # 限制GPU内存使用,减少CUDA冲突
|
|
|
|
|
+ os.environ["FLAGS_fraction_of_gpu_memory_to_use"] = "0.6"
|
|
|
|
|
+
|
|
|
|
|
+ # 使用确定性算法
|
|
|
|
|
+ os.environ["FLAGS_cudnn_deterministic"] = "1"
|
|
|
|
|
+
|
|
|
|
|
+ # 立即释放内存
|
|
|
|
|
+ os.environ["FLAGS_eager_delete_tensor_gb"] = "0.0"
|
|
|
|
|
+
|
|
|
pipeline = create_pipeline(pipeline_name_or_config_path, device=device)
|
|
pipeline = create_pipeline(pipeline_name_or_config_path, device=device)
|
|
|
|
|
|
|
|
should_end = False
|
|
should_end = False
|