Python tf.data.experimental.service.WorkerServer用法及代码示例

in-process tf.data 服务工作者服务器。

用法

tf.data.experimental.service.WorkerServer(
    config, start=True
)

tf.data.experimental.service.WorkerServer 对用户定义的数据集执行 tf.data.Dataset 处理，并通过 RPC 提供结果元素。工作人员与单个 tf.data.experimental.service.DispatchServer 相关联。

dispatcher = tf.data.experimental.service.DispatchServer()
dispatcher_address = dispatcher.target.split("://")[1]
worker = tf.data.experimental.service.WorkerServer(
    tf.data.experimental.service.WorkerConfig(
        dispatcher_address=dispatcher_address))
dataset = tf.data.Dataset.range(10)
dataset = dataset.apply(tf.data.experimental.service.distribute(
    processing_mode="parallel_epochs", service=dispatcher.target))
print(list(dataset.as_numpy_iterator()))
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

启动专用 tf.data 工作进程时，在启动服务器后使用join() 无限期阻塞。

worker = tf.data.experimental.service.WorkerServer(
    port=5051, dispatcher_address="localhost:5050")
worker.join()