浏览代码

chore(zhch): 调整模型推理代码

- 移除限制 GPU 内存使用的环境变量设置
- 注释掉使用确定性算法和立即释放内存的环境变量设置
- 在 worker 函数中添加 paddle 导入和运行检查
zhch158_admin 3 月之前
父节点
当前提交
269880b2ca
共有 1 个文件被更改,包括 6 次插入3 次删除
  1. 6 3
      zhch/mp_infer.py

+ 6 - 3
zhch/mp_infer.py

@@ -12,14 +12,17 @@ from paddlex.utils.device import constr_device, parse_device
 
 
 def worker(pipeline_name_or_config_path, device, task_queue, batch_size, output_dir):
+    import paddle
+    paddle.utils.run_check()
+
     # 限制GPU内存使用,减少CUDA冲突
-    os.environ["FLAGS_fraction_of_gpu_memory_to_use"] = "0.6"
+    # os.environ["FLAGS_fraction_of_gpu_memory_to_use"] = "0.6"
 
     # 使用确定性算法
-    os.environ["FLAGS_cudnn_deterministic"] = "1"
+    # os.environ["FLAGS_cudnn_deterministic"] = "1"
 
     # 立即释放内存
-    os.environ["FLAGS_eager_delete_tensor_gb"] = "0.0"
+    # os.environ["FLAGS_eager_delete_tensor_gb"] = "0.0"
 
     pipeline = create_pipeline(pipeline_name_or_config_path, device=device)