often 5 месяцев назад
Родитель
Сommit
8aa090f0fe

+ 4 - 2
recommend-model-produce/src/main/python/tools/static_ps_infer_v2.py

@@ -65,7 +65,7 @@ class InferenceFetchHandler(object):
         output_dir = os.path.dirname(output_file)
         if not os.path.exists(output_dir):
             os.makedirs(output_dir)
-        
+        self.var_dict = {}  # 用于存储需要获取的变量
         # 创建或清空输出文件
         with open(self.output_file, 'w') as f:
             f.write('')
@@ -107,7 +107,9 @@ class InferenceFetchHandler(object):
         """停止处理"""
         self.terminal_event.set()
         self.finish()
-
+    def set_var_dict(self, var_dict):
+        """设置需要获取的变量字典"""
+        self.var_dict = var_dict
     def wait(self):
         """等待处理完成"""
         self.done_event.wait()