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

TensorFlow2.0 GPU内存管理

最新推荐文章于 2021-08-25 10:46:38 发布

原创最新推荐文章于 2021-08-25 10:46:38 发布 · 1.1k 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #tensorflow #神经网络

小白的ai学习之路同时被 2 个专栏收录

83 篇文章

订阅专栏

TF2

37 篇文章

订阅专栏

本文解决了在使用TensorFlow2.0进行深度学习时遇到的Couldnotcreatecudnnhandle:CUDNN_STATUS_ALLOC_FAILED错误。通过调整GPU内存增长设置，避免了在调用Keras卷积层时出现的内存分配失败问题。

部署运行你感兴趣的模型镜像

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)