当与pyzmq一起使用时，Tensorflow推理挂起

编程问答 2022-07-25

问题描述

我正在尝试使用pyzmq将Tensorflow模型所做的预测传递给另一个进程。

我正在调用像self.model.predict(X)这样简单的代码，其中self.model是Tensorflow模型。

当我单独呼叫此行时，它工作正常。但是，当我使用客户机/服务器体系结构时，客户机和服务器都位于单独的进程中，并连接到与以下相同的套接字端口：

客户：

context = zmq.Context()
socket = context.socket(zmq.REQ)
socket.connect(f'tcp://localhost:1234')
...
socket.send_pyobj(X)
...

服务器：

self.model = tf.keras.models.load_model(...)
context = zmq.Context()
socket = context.socket(zmq.REP)
socket.bind(f'tcp://*:1234')
while True:
    X = socket.recv_pyobj()
    ...
    self.model.predict(X) # <----- HANGS HERE

然后此行挂起。

我注意到的另一件事是，如果在挂起进程时杀死进程，则该端口仍在使用中。

我试图弄清楚如何调试它，因为我似乎无法弄清楚发生了什么。

解决方法

问题是tf.keras.models.load_model(...)和socket.bind(f'tcp://*:1234')的顺序

需要成为

  1. socket.bind(f'tcp://*:1234')
  2. tf.keras.models.load_model(...)

而不是其他方式。

pyzmqtensorflowtensorflowtensorflow