正在运行Tensorflow.Keras模型.仅在一个CPU上预测

Question

我有一个有60个CPU的系统。我想在60个核上并行应用Keras神经网络模型预测。我应该如何强制每个并行进程只使用60个核心中的1个？

from pandarallel import pandarallel
pandarallel.initialize(nb_workers=60)
def my_func(path):
# probably something should be added here to restrict tensorflow.keras model.predict to only one CPU
return my_model.predict(load_and_preprocess(path))
df['prediction'] = df.parallel_apply(lambda x: my_func(x['image_path']))

目前的问题是，对于长度为10的DataFrame，这段代码将永远不间断地运行，而它将在10秒内完成。

Answer 1

为了社区的利益，在此处添加@H4iku答案

import tensorflow as tf
import numpy as np
from multiprocessing import Pool

def _apply_df(data):
model = tf.keras.models.load_model("my_fashion_mnist_model.h5")
return model.predict(data)

def apply_by_multiprocessing(data, workers):
pool = Pool(processes=workers)
result = pool.map(_apply_df, np.array_split(data, workers))
pool.close()
return list(result)

def main():
fashion_mnist = tf.keras.datasets.fashion_mnist
_, (test_images, test_labels) = fashion_mnist.load_data()
test_images = test_images / 255.0
results = apply_by_multiprocessing(test_images, workers=3)
print(test_images.shape)           # (10000, 28, 28)
print(len(results))                # 3
print([x.shape for x in results])  # [(3334, 10), (3333, 10), (3333, 10)]

if __name__ == "__main__":
main()