问题背景
使用tensorflow的 Conv2D 层时,指定了 data_format = “channel_first”,但是如果这时 TF 使用的是CPU,就会报错。
报错如下:
UnknownError: Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above. [Op:Conv2D]
改正方法
在开头加上如下代码即可:
import os
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
该bug的更多解决方法点这里