Возможно ли ускорить обработку hdf5 файла с помощью threading Python 3.6

Имеется список hdf5 файлов (147 штук). В каждом файле хранится набор данных измерений (картинки с интенсивностями, около 2.1М штук), каждая картинка может быть либо slab формата (8192, 128), либо vds формата (16, 512, 128). Т.е. общий формат данных будет выглядеть либо (2.1М, 8192, 128), либо (2.1М, 16, 512, 128) и всё это numpay.array. Необходимо получить одну картинку, которая будет являться суммой интенсивностей со всех файлов.

Хотелось бы узнать у тех, кто имеет явно больший опыт работы с модулями multiprocessing и threading, имеется ли возможность существующий ниже код переписать так, чтобы увеличить скорость обработки данных?

#!/usr/bin/env python3
# coding: utf8
'''
Загрузка основных пакетов
'''
import os
import sys
import h5py as h5 # модуль для работы с hdf5 файлами
import threading
import numpy as np
import subprocess
import re
import argparse

from itertools import product

chunk_size = 12000 # количество картинок, загружаемых в память, далее чанк


os.nice(0)

class CustomFormatter(argparse.RawDescriptionHelpFormatter,
                      argparse.ArgumentDefaultsHelpFormatter):
    pass

def parse_cmdline_args(): # функция для парсинга ключей командной строки
    parser = argparse.ArgumentParser(
        description=sys.modules[__name__].__doc__,
        formatter_class=CustomFormatter)
    parser.add_argument('-i', '--input', nargs='+', type=str, help="List of cxi files") # если на вход подается список файлов для обработки
    parser.add_argument('-f', '--file', type=str, help="File with the list of cxi files") # если на вход подается файл, содержащий список hdf5 файлов для обработки
    parser.add_argument('-p', type=str, help="hdf5 path for the cxi file data") # путь к данным в каждом файле (предполагается, что у всех файлов путь к данным одинаковый)
    parser.add_argument('-m', type=str, help="hdf5 path for the cxi file mask") # путь к маске в каждом файле (предполагается, что у всех файлов путь к маске одинаковый)
    parser.add_argument('-out-prefix', '--outPrefix', type=str, help="Prefix for the output h5 file")  # префикс для выходных файлов с обработанными данными
    return parser.parse_args()


def _read_chunk_at(index_chunk, column_index, data_file):
    global intensity  # матрица суммарной интенсивности
    global data_file_shape  # формат данных

    print('Calculating intensity by {}'.format(threading.currentThread().getName()))

    """
    Чтобы не использовать lock суммируем интенсивности по slice by column,
    результат записываем в соответствующую колонку intensity.
    если, например, у нас имеются следующие данные
    >>a3 = np.array([[[10, 11, 12], [13, 14, 15], [16, 17, 18]],
           [[20, 21, 22], [23, 24, 25], [26, 27, 28]],
           [[30, 31, 32], [33, 34, 35], [36, 37, 38]]])

    >>print(a3)
    array([[[10, 11, 12],
            [13, 14, 15],
            [16, 17, 18]],

           [[20, 21, 22],
            [23, 24, 25],
            [26, 27, 28]],

           [[30, 31, 32],
            [33, 34, 35],
            [36, 37, 38]]])

    то результат должен быть следующий
    np.array([[60, 63, 66], [69, 72, 75], [78, 81, 84])
    который получится следующим образом:
    >>> b = np.zeros((3,3))
    >>> for i in range(3):
    ...     b[:,i] += np.sum(a3[:,:,i], axis=0)
    >>> b
        array([[ 60.,  63.,  66.],
               [ 69.,  72.,  75.],
               [ 78.,  81.,  84.]])

    """
    if len(data_file_shape) > 3:  # проверка на формат данных
         # vds формат (2.1М, 16, 512, 128)
        data = data_file[index_chunk : index_chunk + chunk_size,:,:,column_index]
        intensity[:,:,column_index] += np.sum(data, axis=0)

    else:
         # slab формат (2.1М, 8192, 128)
        data = data_file[index_chunk : index_chunk + chunk_size,:,column_index]
        intensity[:,column_index] += np.sum(data, axis=0)



def main(data_file): 
    # главная функция создания потоков для чтения одного hdf5 файла
    global intensity 
    global data_file_shape
    global length

    print('Number of patterns in file is {}'.format(length))
    print('Creating Threads')

    args_iter = product(range(0, length, chunk_size),range(0, data_file_shape[-1])) # индекс для начала считывания чанка, индекс ряда, куда будет записывать поток просуммированную интенсивность чанка
    
    threads = list()

    for index in args_iter:
        # creating threads
        thread = threading.Thread(target=_read_chunk_at, args=(index[0], index[1], data_file))
        threads.append(thread)
        # start threads
        thread.start()

    for index, thread in enumerate(threads):
        # wait until threads complete
        thread.join()
        print("Main    : thread {} done".format(index))


if __name__ == "__main__":

    args = parse_cmdline_args()

    list_of_cxi_files = []

    if args.input is not None:
        list_of_cxi_files = args.input
    else:
        try:
            with open(args.file, 'r') as f:
                for line in f:
                    list_of_cxi_files.append(line.strip())
        except:
            print("It's necessary to give a list of files or stream with them as input parameter.\n")


    data_label = args.p
    mask_label = args.m

    total_num_patterns = 0

    data_file_shape = subprocess.check_output(['/opt/hdf5/hdf5-1.10.6/bin/h5ls', str(list_of_cxi_files[0])+str(data_label)])
    data_file_shape = data_file_shape.strip().decode('utf-8').split('Dataset ')[1]
    data_file_shape = re.sub(r'({|}|/Inf)', '', data_file_shape).split(', ')

    data_file_shape = [int(i) for i in data_file_shape]

    print('Shape of data in files is {}'.format(data_file_shape[1:]))

    intensity = np.zeros(data_file_shape[1:])

    print('It is necessary to process {} files'.format(len(list_of_cxi_files)))

    for filename in list_of_cxi_files:
        print('{} is processing'.format(filename))
        h_file = h5.File(filename, 'r')
        data_file = h_file[data_label]
        length = data_file.shape[0]
        total_num_patterns += length
        main(data_file)
        h_file.close()
        print('\n')
    
    print('Total number of patterns is {}'.format(total_num_patterns))

    if args.outPrefix is None:
        f_name = 'parallel-av-num-patterns.h5'
        f2_name = 'parallel-av-count-dots.h5'
        f3_name = 'count-dots-data.h5'
    else:
        f_name = args.outPrefix + '-av-num-patterns.h5'
        f2_name = args.outPrefix + '-av-count-dots.h5'
        f3_name = args.outPrefix + '-count-dots-data.h5'
    
    f = h5.File(f_name, 'w')
    f.create_dataset('/data/data', data=np.array([intensity / total_num_patterns]))
    f.close()

    print('Finish\n')

global intensity - переменная, в которой хранится усредненная итоговая интенсивность со всех картинок из всех файлов.


Ответы (1 шт):

Автор решения: eri

Тут лучше подойдёт multiprocessing.Pool().imap_unordered()

На вход подаешь имя файла, на выход интенсивность в этом файле.

Из мапа выйдет итератор с интенсивностями, который загоняешь в reduce(). Эту операцию можно тоже распараллелить, но тут особо смысла нет потому как простое суммирование.

→ Ссылка