1. Problem description
When using tensorflow-gpu, the following error occurred:
tensorflow.python.framework.errors_impl.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.
[[node sequential/conv2d/Conv2D (defined at tf_keras_classification_model-cnn.py:107) ]] [Op:__inference_distributed_function_930]
Function call stack:
distributed_function
Second, the solution
Using tensorflow, if you don't add settings, even a small model will occupy the entire GPU, resulting in a waste of resources. So we need to set up so that the program uses the GPU on demand.
Add the following lines of code before the code to implement tensorflow's on-demand GPU allocation
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
# 定义TensorFlow配置
config = ConfigProto()
# 配置GPU内存分配方式，按需增长，很关键
config.gpu_options.allow_growth = True
# 在创建session的时候把config作为参数传进去
session = InteractiveSession(config=config)