Question

我使用tensorflow为卷积神经网络构建了一个输入管道。特别是，在评估时，我需要从磁盘读取一个图像并从中生成几个（可能是100000个）补丁。这是使用 tf.extract_image_patches 完成的。然后，我创建批次（每个2000个补丁），然后将它们提供给网络。

问题是该程序分配了补丁的整个张量，导致我的4GB GPU出现Out Of Memory问题。我希望它只需要在每一步分配评估中涉及的批次。有没有解决方法来避免这种情况？

错误的

OOM when allocating tensor with shape[104720,50,50,3]

     [[Node: Cast = Cast[DstT=DT_FLOAT, SrcT=DT_UINT8, _device="/job:localhost/replica:0/task:0/device:GPU:0"](MirrorPad/_63)]]

     [[Node: Gather/_65 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_22_Gather", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

最小代码

imagepath = tf.convert_to_tensor(imagepath, dtype=tf.string)
image= tf.train.slice_input_producer(imagepath,shuffle=False,num_epochs=1)
image = tf.read_file(image)
image = tf.image.decode_jpeg(image, channels=3)
image=tf.extract_image_patches(image, ksizes=[1,50, 50,1], strides=[1,25,25,1], rates=[1,1,1,1],padding="VALID")
image=tf.reshape(image,(-1,50,50,3))

X = tf.train.batch([image], batch_size=2000,
                            capacity=batch_size * 8,
                                num_threads=1,
                                allow_smaller_final_batch=True,
                                enqueue_many=True
                                )

调用tf.extract_image_patches时的Memoy

0 个答案: