python - TensorFlow 未使用所有系统内存

Question

我正在尝试使用 tensorflow 数据集 API 加载相当大的 h5 文件。

我最大的张量是大约 800MB 的内存（复值32x15x640x322）。

加载它时，张量流给了我以下警告：

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

然而，我的系统有 251GB 的 RAM（来自htop），所以 800MB 远不及系统内存的 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、tf 2.1 完成的。

标签： pythontensorflowmemory