Question

问题

我正在按照here准则构建自定义TensorFlow GPU操作。一些操作将许多张量用作输入，我担心这些张量的组合大小太大而无法在某些GPU上容纳。

设置

我正在运行带有TF 1.8（带有GPU支持）的Ubuntu 18.04，并使用Python 2.7。 GPU是NVIDIA Quadro P400，我有CUDA Toolkit 9.0和cuDNN SDK v7。

当前，我构建一个TensorFlow图，向其添加自定义操作，导出.pb文件，然后对生成的图进行一些简单的图像分类。我想发生的事情是在构建图形时知道给定的操作是否会导致GPU内存不足。如果没有，那么我将合并操作。另一方面，如果我看到该操作没有内存，那么我将实现另一版本的操作，该版本速度较慢但使用的内存较少。

我尚未在P400上遇到内存问题，但是理想情况下，无论内存是否可用，这都会推广到任何GPU。为了模拟这一点，我使用了

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.05
with tf.Session(config=config, ...) as sess:
    ...

严格限制TF可用的内存量。我了解到这种情况在实践中可能不会经常发生，但我假设我的自定义操作的内存需求会任意增大。

在测试自定义操作时，我已经能够对其进行足够精确的调整，以使一个实现可以成功运行，而另一个使用稍多的内存则会导致OOM错误。通过在report_tensor_allocations_upon_oom = True中设置tf.RunOptions，可以看到将哪些张量分配给了哪些设备。典型输出如下：

ResourceExhaustedError (see above for traceback): OOM when allocating tensor with shape[512,512,3,3] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
     [[Node: resnet_v2_50/block4/unit_1/bottleneck_v2/conv2/Conv2D = Conv2D[T=DT_FLOAT, data_format="NCHW", dilations=[1, 1, 1, 1], padding="SAME", strides=[1, 1, 1, 1], use_cudnn_on_gpu=true, _device="/job:localhost/replica:0/task:0/device:GPU:0"](resnet_v2_50/block4/unit_1/bottleneck_v2/conv1/Relu, resnet_v2_50/block4/unit_1/bottleneck_v2/conv2/weights)]]

Current usage from device: /job:localhost/replica:0/task:0/device:GPU:0, allocator: GPU_0_bfc
  392.0KiB from resnet_v2_50/block4/unit_1/bottleneck_v2/shortcut/Conv2D
  98.0KiB from resnet_v2_50/block4/unit_1/bottleneck_v2/conv1/Conv2D

     [[Node: TopKV2/_13 = _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_783_TopKV2", tensor_type=DT_INT32, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

Current usage from device: /job:localhost/replica:0/task:0/device:GPU:0, allocator: GPU_0_bfc
  392.0KiB from resnet_v2_50/block4/unit_1/bottleneck_v2/shortcut/Conv2D

最让我感到困惑的是，这似乎不是浪费资源的自定义操作，而是一种卷积。但是，当我在没有自定义操作（或使用需要较少内存的其他操作）的情况下运行时，没有错误。

我的尝试

我已经尝试过config.gpu_options.allow_growth = True，但发现它不适合我的需求。我还使用了here定义的函数BytesInUse()，BytesLimit()和MaxBytesInUse()，但是它们同样无济于事。

我的一个想法是在没有任何自定义操作的情况下运行网络。然后，我可以将BytesLimit() - MaxBytesInUse()视为自定义操作的“预算”。到目前为止，这还行不通，因为TensorFlow会根据可用性自动增加其内存消耗。

目前，我唯一可行的解决方案是缓慢增加自定义操作的大小，直到它们过大并给出我上面提到的错误为止。这显然是不理想的。

如果有某种方法可以强制TensorFlow使用可能的绝对最小内存，那么只需运行一次网络并测量资源消耗即可。

任何帮助或想法将不胜感激。

如何确保TensorFlow自定义操作将适合GPU内存？

问题

设置

我的尝试

0 个答案: