Сохранение в файл hdf5 с использованием h5py замедляется со временем - PullRequest
Новая
       9

Сохранение в файл hdf5 с использованием h5py замедляется со временем

0 голосов
/ 29 августа 2018

Я хочу хранить> 1M изображений в файле hdf5. Проблема сохранения их в цикле for со временем замедляется. Он начинается с 80 часов для 1,5 М изображений и уменьшается до 130 часов через полчаса после обработки изображений 6K. Данные поступают в виде списка кортежей. Для каждого изображения есть кортеж, содержащий идентификатор и URL. Функция download_image проверяет, есть ли идентификатор в hdf5 id_dset, и загружает изображение. Функция save_image увеличивает hdf5 image_dset и hdf5 id_dset, если это необходимо, и сохраняет id и изображение. В конце концов, все это идет к многопроцессорному пулу. 

def download_image(id_url):    
    (id, url) = id_url

    if id in id_dset:
        log_file.write('Image {} already exists. Skipping download\n'.format(id))
        return None  

    try:
        response = request.urlopen(url)
        image_data = response.read()
    except:
        log_file.write('Could not download image {} from {}\n'.format(id, url))
        return None

    try:
        image = Image.open(BytesIO(image_data))        
    except:
        log_file.write('Failed to parse image {}\n'.format(id))
        return None

    try:
        image = image.convert('RGB')
    except:
        log_file.write('Failed to convert image {} to RGB\n'.format(id))
        return None

    return id, image

def save_image(idx, id, image):
    try:
        if image_dset.shape[0] <= idx:
            image_dset.resize(idx+1000, axis=0)
            id_dset.resize(idx+1000, axis=0)

        image_dset[idx] = np.array(image)
        id_dset[idx] = id

    except:
        log_file.write('Failed to save image {}\n'.format(id))

# Create an empty file
chunk_size = 10
input_dir = '../input/images-h5/'

if not os.path.exists(input_dir):
    os.makedirs(input_dir)

with h5py.File(input_dir + 'images.h5', 'w') as f:

    image_dset = f.create_dataset(name='image', shape=(0, 400, 400, 3), maxshape=(None, 400, 400, 3), dtype=np.uint8, 
                               chunks=(chunk_size, 400, 400, 3), compression='gzip', compression_opts=3)

    dt = h5py.special_dtype(vlen=str)
    id_dset = f.create_dataset(name='id', shape=(0,), maxshape=(None,), dtype=dt)

from multiprocessing.pool import ThreadPool
pool = ThreadPool(6)

log_file = open(input_dir + 'log_{}.txt'.format(datetime.datetime.now().isoformat()), 'w')

# Add to file
with h5py.File(input_dir + 'images.h5', 'r+', libver='latest') as f:

    image_dset = f['image']
    id_dset = f['id']

    if id_dset.shape[0] == 0:
        idx = 0        
    else: 
        idx = np.argmax(id_dset[:] == '')

    for res in tqdm.tqdm_notebook(pool.imap(download_image, id_url_list), total=len(id_url_list)):
        if res is not None:
            id, image = res

            save_image(idx, id, image)   

            idx += 1

log_file.close()
...