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)
发布了62 篇原创文章 · 获赞 33 · 访问量 3471

猜你喜欢

转载自blog.csdn.net/python_LC_nohtyp/article/details/104364071
今日推荐