Python 类方法调用上的Tensorflow集合图/会话上下文

我正在并行运行两个tensorflow模型。每个都有自己的图表和会话

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

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

在通话中，我使用

和
设置上下文：

    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装饰器将前两行包装在方法中，如predict方法，如为变量范围所述。不过，这将是相同的代码，只是外观更好，并且在有多个这样的方法时非常方便

我不知道你是否会认为它更优雅，但我个人更喜欢直接调用session对象，而不是使用自己的<代码> .sas.asyDebug（）：控制流结构。