python之Tensorflow 在类方法调用上设置图形/ session 上下文

我并行运行两个 tensorflow 模型。每个都有自己的图表和 session 。

为了确保没有图形/操作冲突，我创建了一个包装类来维护两个模型各自的图形和 session :

class ModelWrapper: 
    def __init__(self): 
        self.graph = None 
        self.sess = None 
        self.model = None 
 
    def load_model(self, pth_model=None): 
        # load graph, session, and model here 

在通话中，我然后使用 with 设置上下文:

    def predict(self, np_x): 
        with self.graph.as_default(): 
            with self.sess.as_default(): 
                np_y_preds = self.model.predict(np_x) 
        return np_y_preds 

每次调用predict方法时，是否有更优雅的方式来设置上下文？

请您参考如下方法:

您可以使用 python 装饰器将前两行包装在方法中，就像您描述的预测方法一样 here对于变量范围。不过，那将是相同的代码，只是更好看，并且在有多个这样的方法时非常方便。

我不知道你是否会认为这更优雅，但我个人更喜欢直接调用 session 对象，而不是使用 with self.sess.as_default():控制流结构。