Загрузка модели из tenorflow SavedModel на несколько графических процессоров

Question

Допустим, кто-то вручил мне TF SavedModel, и я хотел бы воспроизвести эту модель на 4 графических процессорах, установленных на моей машине, чтобы я мог параллельно выполнять вывод для пакетов данных.Есть ли хорошие примеры того, как это сделать?

Я могу загрузить сохраненную модель следующим образом:

def load_model(self, saved_model_dirpath):
    '''Loads a model from a saved model directory - this should 
       contain a .pb file and a variables directory'''

    signature_key = tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
    input_key = 'input'
    output_key = 'output'

    meta_graph_def = tf.saved_model.loader.load(self.sess, [tf.saved_model.tag_constants.SERVING],
                                                saved_model_dirpath)
    signature = meta_graph_def.signature_def

    input_tensor_name = signature[signature_key].inputs[input_key].name
    output_tensor_name = signature[signature_key].outputs[output_key].name

    self.input_tensor = self.sess.graph.get_tensor_by_name(input_tensor_name)
    self.output_tensor = self.sess.graph.get_tensor_by_name(output_tensor_name)

.. но для этого потребуется дескриптор сеанса.Для моделей, которые я написал сам, у меня был бы доступ к функции вывода, и я мог бы просто вызвать ее и обернуть, используя with tf.device(), но в этом случае я не уверен, как извлечь функцию вывода из сохраненногоМодель.Должен ли я загрузить 4 отдельных сеанса или есть лучший способ?Не могу найти много документации по этому вопросу, но заранее извиняюсь, если что-то пропустил.Спасибо!

Answer 1

В данный момент в TensorFlow нет поддержки этого варианта использования.К сожалению, «репликация функции вывода», основанной только на SavedModel (который является в основном графом вычислений с некоторыми метаданными), является довольно сложной (и хрупкой, если реализуется) проблемой преобразования графа.

Если вы неЧтобы иметь доступ к исходному коду, который создал эту модель, лучше всего загрузить SavedModel 4 раза в 4 отдельных графика, каждый раз переписывая целевое устройство в соответствующий графический процессор.Затем запустите каждый график / сеанс отдельно.

Обратите внимание, что вы можете вызывать sess.run() несколько раз одновременно, поскольку sess.run() освобождает GIL для времени фактических вычислений.Все, что вам нужно, это несколько потоков Python.