Tensorflow在类方法调用上设置图/会话上下文

Question

我正在并行运行两个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

随叫随到我然后使用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方法时，是否有更优雅的方法来设置上下文？

Answer 1

你可以使用python修饰器将前两行包装在像预测方法这样的变量范围的here之类的方法中。那将是相同的代码，只要更好看，并且在拥有多个这样的方法时非常方便。

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