tensorflow和keras-从SQS进行生产中的快速推理



我用Keras编写和训练了我的模型。我正试图在生产中使用它进行推理。我收到包含(path_in, path_out)元组的SQS"任务"消息。

我显然可以使用:

BATCH_SIZE = 10
batch_messages = []

while True:
while len(batch_messages) < BATCH_SIZE:
msg = sqs.read_messsage()  
batch_messages.apend(msg)
assert len(batch_messages) == BATCH_SIZE
batch = np.array([read_image(msg.path_in) for msg in batch_messages]) 
output_batch = model.predict(batch)
for i in range(BATCH_SIZE):
write_output(output_batch[i], path=batch_messages[i].path_out)
batch_messages = []

这样做的问题是,代码浪费了从SQS读取、从磁盘读取图像并在最后将其写回的大部分时间。这意味着GPU在此期间处于空闲状态。

我知道Keras的Sequence,但不确定它是否也适用于这种情况,以及用于推断(而不是训练(

我建议您使用Tensorflow Serving解决方案,因为它实现了服务器端的批处理策略,可以优化推理速度和GOU利用率。此外,如果你想加快你的管道,你应该将模型转换为TensorRT模型,该模型可以优化特定GPU的模型操作(它做得更多(。

最新更新