CODEWITHSUNDEEP
pandassessionmultiprocessingtf.keraspredict
Vahid the Great
Vahid the Great

Reputation: 483

Running Tensorflow.Keras Model.Predict on Only One CPU

I have a system with 60 CPUs. I want to apply a Keras Neural Net model prediction on 60 cores in parallel. How should I force each parallel process to use only 1 of the 60 cores?

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']))

The problem is that at the moment, this code runs non-stop forever while it finishes in 10 seconds for a DataFrame of length 10.

Upvotes: 2

Views: 699

Answers (1)

user11530462
user11530462

Reputation:

For the benefit of community adding @H4iku answer here

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()

Upvotes: 1

Related Questions