Как подружить Keras и Multiprocessing?

Question

[Evgeny_A]

Здравствуйте.

У меня база в которой хранятся картинки. Есть нейронка, собранная и сохраненная с помощью Keras. Скрипт должен получить строки из таблицы и по порядку обработать с помощью нейронки изображение в каждой из них, а результат обработки вернуть в строку. Чтобы это не занимало слишком много времени, сразу решил использовать метод Pool из пакета Multiprocessing. Вот в такой реализации, если pools_number > 1 и используется multiprocessing, скрипт виснет на keras_model.predict([x]). Если pools_number = 1, то все работает как надо:

pools_number = 1

from keras.models import load_model
from keras.preprocessing import image as keras_image

# Загружаем сохраненную сеточку в память
keras_model = load_model('/work_model.h5')

def getResult(image):

	print('Обрабатываем {}...'.format(image['name']))

	start = time.monotonic()

	# Открываем картинку
	img_keras = keras_image.load_img(image['path'], target_size=(150, 150))
	# Конвертируем в np массив
	x = keras_image.img_to_array(face_img)
	# Добавляем ось
	x = np.expand_dims(x, axis=0)  
	x /= 255.

	# Получаем результат
	predict_result = keras_model.predict([x])
	result1 = predict_result[0][0]
	result2 = predict_result[0][1]

	result_time = time.monotonic() - start
	print('Выполнено за {:>.3f} секунд'.format(result_time))

rows = [image1, image2]

if pools_number == 1:

	for idx, row in enumerate(rows):

		getResult(row)

else:

	with multiprocessing.Pool(pools_number) as pool:

		pool.map(getResult, rows)
		pool.close()
		pool.join()

Читал статейки и накидал второй вариант, когда модель сеточки из файла загружается внутри функции getResult():

pools_number = 1

def getResult(image):

	from keras.models import load_model
	from keras.preprocessing import image as keras_image

	# Загружаем сохраненную сеточку в память
	keras_model = load_model('/work_model.h5')

	print('Обрабатываем {}...'.format(image['name']))

	start = time.monotonic()

	# Открываем картинку
	img_keras = keras_image.load_img(image['path'], target_size=(150, 150))
	# Конвертируем в np массив
	x = keras_image.img_to_array(face_img)
	# Добавляем ось
	x = np.expand_dims(x, axis=0)  
	x /= 255.

	# Получаем результат
	predict_result = keras_model.predict([x])
	result1 = predict_result[0][0]
	result2 = predict_result[0][1]

	result_time = time.monotonic() - start
	print('Выполнено за {:>.3f} секунд'.format(result_time))

rows = [image1, image2]

if pools_number == 1:

	for idx, row in enumerate(rows):

		getResult(row)

else:

	with multiprocessing.Pool(pools_number) as pool:

		pool.map(getResult, rows)
		pool.close()
		pool.join()

В этом случае все работает замечательно при любых значениях pools_number, но так как файл work_model.h5 весит почти гигабайт, то обработка одного изображения занимает почти 30 секунд, а требования к объему памяти сервера (чтобы поддерживать, например, 10-14 потоков для 12 ядерного процессора) взлетают до небес. Отказаться от использования многопоточности не могу, так строк в базе данных больше 8 миллионов, а время на выполнение задачи ограничено.

Вопрос: как подружить Keras и Multiprocessing?

Answer 1

[Andrey Dugin]

0) Используйте GPU.
1) Используйте батчи, т.е. обрабатывайте изображения пачкой.
2) Процедуры загрузки-предобработки изображений необходимо измерить по времени исполнения и заменить на быстрые - например, из OpenCV (cv2.imread, cv2.warpAffine, cv2.normalize). Также заранее подготовить буферы для результата, а не аллоцировать память под выходной массив каждый раз заново.
3) Загрузку изображений можно вынести в асинхронный код.
4) Tensorflow (или что там у вас под капотом) и так использует все ядра.
5) Вы оптимизировали нейронку для инференса? Frozen graph, merge BN layers, TensorRT, снижение разрядности?
6) Почему модель весит почти гигабайт? Более лёгкие модели не подошли? Google MobileNetV2, например?

Comments

[Evgeny_A]

Спасибо за ответ. Одна мысль стала решением моей задачи (2 пункт). По другим пунктам:

1. GPU нет, действую в условиях ограниченных ресурсов
   
2. Это помогло, спасибо!
   
3. В моем случае самым долгим местом была предсказание Keras
   
4. Почти так и поступил)
   
5. Не знал, спасибо
   
6. Буду учить. В рамках текущей задачи нужно было просто решение на уровне циклов.
   
7. Потому что такой она получилась после обучения и в оптимизации я еще ничего не понимаю, но применять уже нужно было)
   

Answer 2

[Evgeny_A]

В общем Keras устроен так, что единственный способ совместить его с Multiprocessing, это загружать модель отдельно для каждого процесса. Если процессы создаются для каждого "мелкого" кода, как, например, в моем случае для обработки одного изображения, то можно перекроить скрипт так, чтобы процесс создавался не для отдельной картинки, а я для их группы (в моем случае я одновременно 20 картинок загружал). Таким образом требования к памяти станут лояльнее и можно уложить в имеющиеся ресурсы.