Использование nondifferentiable_batch_function не выполняет операции на GPU при выводе с использованием Tensorflow Serving - PullRequest
Новая
       57

Использование nondifferentiable_batch_function не выполняет операции на GPU при выводе с использованием Tensorflow Serving

0 голосов
/ 06 марта 2020

Я пытаюсь проверить возможность использования nondifferentiable_batch_function для вывода. В частности, мне нужно собирать пакеты и выполнять их на GPU для лучшей пропускной способности. Я знаю, что есть некоторые другие подходы (например, использование пакетного обслуживания tenorflow), но использование функции пакетной обработки в графике может быть полезным для моей текущей настройки. Экспериментируя с различными подходами, я обнаружил, что размещение устройства для операций, упакованных в пакетную функцию, может быть несовместимым. Мне удалось выполнить операции на графическом процессоре, включив конкретную функцию в пакетную функцию, а затем сохранить модель. Тем не менее, когда я пытаюсь использовать сохраненную модель через тензор потока, она все равно выполняет операции с процессором. Вот как я создал сохраненную модель: 

signature = (tf.TensorSpec(shape=[None, 3], dtype=tf.float32, name="x"),)

@tf.function(input_signature=signature)
def run(x):
  with tf.device("/GPU:0"):
        return tf.nn.relu(tf.matmul(x, tf.random.uniform((3, 1))))

concrete_run = run.get_concrete_function()

@tf.nondifferentiable_batch_function(1, 6, 10)
def batch_run(x):
    return concrete_run(x)

class Model(tf.Module):
    @tf.function(input_signature=signature)
    def run(self, x):
        return batch_run(x)

model = Model()
x = tf.random.uniform((3, 3))

# MatMul and Relu are executed on CPU
model.run(x)

concrete_model_run = model.run.get_concrete_function()

# MatMul and Relu are executed on CPU
concrete_model_run(x)

tf.saved_model.save(model, "model", signatures={ "run": model.run })
model = tf.saved_model.load("model")

# MatMul and Relu are executed on GPU
model.signatures["run"](x=x)

Пожалуйста, найдите полный тест там https://colab.research.google.com/drive/1ZcIbnixdcQ2OSabEAlcf5SYYdkoRzYtu

...