tensorflow设置图形/会话方法在类方法调用上

Question

我并行运行两个Tensorflow模型。每个都有自己的图形和会话。

为了确保没有图形/OP冲突，我创建了一个包装程序类，该类别维护了两个模型的相应图形和会话：

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装饰器将前两条线包装在方法中，例如您的预测方法，如此处所述的可变范围。不过，那将是相同的代码，在使用多种此类方法时，它可以更好地看，并且非常方便。

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