应用预取数据集的 dataset.repeat

Question

我正在尝试实施 AdaRound 量化算法，我需要一层一层地训练我的层。

我正在使用 1024 的数据集，批量大小为 32，我需要对数据集进行大约 312 次迭代（或对批处理数据集进行 10k 次迭代），我注意到数据是从主机每次迭代都连接到设备并且数据没有缓存在 GPU 上（尽管重复使用相同的数据）- GPU 有 30%~40% 的时间处于空闲状态

Idle GPU percentage

数据仍然在高级迭代中从主机复制到设备：

memcpyH2D chunk in single iteration

我试过使用

tf.data.experimental.prefetch_to_device

tf.data.experimental.copy_to_device

但是当我在 prefetch_to_device 或 copy_to_device 之后迭代数据时，张量存储在 GPU 上，但是如果我使用 repeat 遍历数据集，张量存储在CPU

我尝试在没有 dataset.repeat 的情况下使用 model.fit 但有多个时期，但我得到了类似的行为。

我还尝试将 model.fit 与存储在 GPU 上的张量一起使用，但模型的拟合将其转换为数据集，这迫使数据返回到 CPU

重现问题的代码片段：

input_shape = (56, 56, 64)
output_shape = (54, 54, 64)
conv = tf.keras.layers.Conv2D(64, (3, 3))
mock_input = tf.keras.layers.Input(input_shape)
mock_output = conv(mock_input)
train_model = tf.keras.Model(inputs=mock_input, outputs=mock_output)

input_data = np.random.rand(1024, *input_shape)
output_data = np.random.rand(1024, *output_shape)

input_dataet = tf.data.Dataset.from_tensor_slices(input_data)
output_dataset = tf.data.Dataset.from_tensor_slices(output_data)

train_model.compile(
    optimizer='adam',
    loss='mse'
)
train_data = tf.data.Dataset.zip((input_dataet, output_dataset))
batched_train_data = train_data.batch(32).cache()
fetched_train_data = batched_train_data.prefetch(tf.data.AUTOTUNE).repeat()
with tf.profiler.experimental.Profile('logs'):
    train_model.fit(fetched_train_data, steps_per_epoch=1024, epochs=1)

有没有办法在 GPU 上应用 dataset.repeat 操作？

我在 python 3.6.9

Answer 1

详细解答

nvidia 有一个名为 Nvidia DALI, this packages offers an efficient wrapper to tensorflow's dataset (and more, but this is the relevant feature I used here), I had to install 2 packages - nvidia-dali-cuda110, nvidia-dali-tf-plugin-cuda110 (a detailed installation guide can be found here)

的软件包

我用过的class叫DALIDataset, to insatiate it properly I first had to initialize pipeline object

single iteration 使用正确初始化的 DALIDataset

代码片段：

from nvidia.dali.plugin.tf import DALIDataset
from nvidia.dali import pipeline_def, fn

def prep_dataset_dali(dir1, dir2, batch_size):
    @pipeline_def(batch_size=batch_size, num_threads=3, device_id=0)
    def pipe(path1, path2):
        data1 = fn.readers.numpy(device='cpu', file_root=path1, file_filter='*.npy')
        data2 = fn.readers.numpy(device='cpu', file_root=path2, file_filter='*.npy')
        return data1.gpu(), data2.gpu()
    my_pipe = pipe(dir1, dir2)
    my_pipe.build()
    return DALIDataset(my_pipe, output_dtypes=(tf.float32, tf.float32), output_shapes=((None, 56, 56, 64), (None, 56, 56, 64)))

注：

外部管道不适用于 DALIDataset，但它可能适用于实验部分DALIDatasetWithInputs class

应用预取数据集的 dataset.repeat

apply dataset.repeat of prefetched dataset

keras

tensorflow

tensorflow-datasets

tf.keras

tensorflow2.0