Tensorflow не использует всю системную память - PullRequest
Новая
       4

Tensorflow не использует всю системную память

0 голосов
/ 05 мая 2020

Я пытаюсь загрузить довольно большие файлы h5 с помощью API набора данных tensorflow.

Мой самый большой тензор занимает около 800 МБ в памяти (комплексное значение 32x15x640x322).

При его загрузке tensorflow дает мне следующее предупреждение: 

2020-05-05 17:41:32.662641: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 791347200 exceeds 10% of system memory.

Однако в моей системе ОЗУ составляет 251 ГБ (из htop), поэтому 800 МБ - это далеко не 10% системной памяти. Я не понимаю, почему появляется это предупреждение.

Многие из текущих решений этой проблемы предлагают уменьшить размер пакета, но это не моя проблема. Я хочу понять, почему у меня появляется это предупреждение, даже если кажется, что у меня достаточно памяти.

Если кто-то хочет воспроизвести это, вам просто нужно установить мой пакет с GitHub (!pip install git+https://github.com/zaccharieramzi/fastmri-reproducible-benchmark). Затем вы можете запустить следующий код: 


import h5py
import numpy as np
import tensorflow as tf
from tqdm.notebook import tqdm

from fastmri_recon.data.utils.h5 import from_multicoil_train_file_to_image_and_kspace_and_contrast

K_shape_single_coil = (32, 640, 322)
K_shape_multi_coil = (32, 15, 640, 322)
I_shape = (32, 320, 320)
contrast = 'CORPD_FBK'

def create_data(filename, multicoil=False):
    k_shape = K_shape_single_coil
    image_ds = "reconstruction_esc"
    if multicoil:
        k_shape = K_shape_multi_coil
        image_ds = "reconstruction_rss"
    kspace = np.random.normal(size=k_shape) + 1j * np.random.normal(size=k_shape)
    image = np.random.normal(size=I_shape)
    kspace = kspace.astype(np.complex64)
    image = image.astype(np.float32)
    with h5py.File(filename, "w") as h5_obj:
        h5_obj.create_dataset("kspace", data=kspace)
        h5_obj.create_dataset(image_ds, data=image)
        h5_obj.attrs['acquisition'] = contrast

n_files = 5
for i in range(n_files):
    create_data(f'test_{i}.h5', multicoil=True)


path = './'


files_ds = tf.data.Dataset.list_files(f'{path}*.h5', seed=0)

res = next(iter(files_ds))

rand = False
parallel = False
inner_slices = None

selection = [
    {'inner_slices': inner_slices, 'rand': rand},  # slice selection
    {'rand': parallel, 'keep_dim': False},  # coil selection
]
def _tf_filename_to_image_and_kspace_and_contrast(filename):
    def _from_train_file_to_image_and_kspace_and_contrast_tensor_to_tensor(filename):
        filename_str = filename.numpy()
        image, kspace, contrast = from_multicoil_train_file_to_image_and_kspace_and_contrast(
            filename_str,
            selection=selection,
        )
        return tf.convert_to_tensor(image), tf.convert_to_tensor(kspace), tf.convert_to_tensor(contrast)
    [image, kspace, contrast] = tf.py_function(
        _from_train_file_to_image_and_kspace_and_contrast_tensor_to_tensor,
        [filename],
        [tf.float32, tf.complex64, tf.string],
    )
    if rand:
        n_slices = (1,)
    else:
        n_slices = (inner_slices,)
    if parallel:
        kspace_size = (640, None)
    else:
        kspace_size = (15, 640, None)
    image_size = (320, 320)
    image.set_shape(n_slices + image_size)
    kspace.set_shape(n_slices + kspace_size)
    return image, kspace, contrast

filename = res.numpy().decode("utf-8")
image, kspace, contrast = _tf_filename_to_image_and_kspace_and_contrast(filename)

Последняя функция, конечно, будет в наборе данных, но я поставил ее для целей этого примера.

Это делается с помощью python 3.6.10, тс 2.1.

...