График Tensorflow Set / контекст сеанса при вызове метода класса

Question

Я использую две модели tenorflow параллельно. У каждого свой график и сессия.

Чтобы гарантировать отсутствие конфликтов графиков / операций, я создал класс-оболочку, поддерживающий соответствующий график и сеанс двух моделей:

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(): структуры потока управления.