Question

我正在使用tf.data.Dataset API构建数据管道，但出现OOM错误。假设我已经拥有features和labels，它们是[N，H，W，C]顺序的4D numpy数组。这是我创建dataset对象的方法：

batch_size = 100
num_samples = features.shape[0] # number of training samples

features_placeholder = tf.placeholder(tf.float32, [None, feature_size], name='features_placeholder')
labels_placeholder = tf.placeholder(tf.float32, [None, label_count], name='labels_placeholder')

dataset = tf.data.Dataset.from_tensor_slices((features_placeholder, labels_placeholder))
dataset = dataset.batch(batch_size)
dataset = dataset.shuffle(num_samples)
dataset = dataset.prefetch(buffer_size=1)
iterator = dataset.make_initializable_iterator()
init_op = iterator.initializer

我使用tf.placeholder的原因可以参考this guide，它基本上建议如果数据是大型numpy数组，则使用dataset来定义tf.placeholder来节省内存。我的训练数据集中有54368个样本）。训练部分如下：

for i in range(epoch):
    sess.run([init_op, optimizer], 
             feed_dict={features_placeholder:features, labels_placeholder:labels]}

但是我看到一条错误消息：

OOM when allocating tensor with shape[54368,40,3,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc

我回溯的

发生在模型中定义的tf.layers.conv2d层上。我该如何解决这个OOM问题？

Answer 1

在shuffle文档中，... fills a buffer with buffer_size elements ...写道，因此，在您的情况下，您的数据集将至少占用54368 * 40 * 3 * 64 * 32 * 2位，大约3.4GB。仅用于随机操作。您是否在使用4GB的GPU？

另一件事是，预取buffer_size应该大于1。为什么要预取1个元素，也许是一两个批次？

使用tf.data.Dataset时如何解决OOM错误？

1 个答案: