Tensorflow2.0 之Could not create cudnn handle: CUDNN_STATUS_ALLOC_FAILED问题

Tensorflow2.0 之Could not create cudnn handle: CUDNN_STATUS_ALLOC_FAILED问题

问题描述：

在tensorflow2.0的学习过程中，遇到了Could not create cudnn handle: CUDNN_STATUS_ALLOC_FAILED，发现这个问题是我在调用kseras模块的卷积类前向传播时触发的，之前的全连接层都没有出现过类似的问题。分析问题应该是gpu没有正确的分配所导致的。

import tensorflow as tf
from tensorflow.keras import layers
x = tf.random.normal([2, 5, 5, 3])
w = tf.random.normal([3, 3, 3, 4])
layer = layers.Conv2D(4, kernel_size=(3, 4), strides=(2, 1), padding='SAME')
out = layer(x)
print(out.shape)

解决方法

解决方法也很简单，只需要在导入模块后加上两行代码即可：

devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(devices[0], True)

加上后代码长这样：

import tensorflow as tf
from tensorflow.keras import layers
devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(devices[0], True)
x = tf.random.normal([2, 5, 5, 3])
w = tf.random.normal([3, 3, 3, 4])
layer = layers.Conv2D(4, kernel_size=(3, 4), strides=(2, 1), padding='SAME')
out = layer(x)
print(out.shape)