单GPU下Jupyter多内核运行TensorFlow报BLAS支持错误如何解决

硬件配置 #

仅有1块GPU，运行设备查询代码的输出如下：

[name: "/device:CPU:0" device_type: "CPU" memory_limit: 268435456
locality {} incarnation: 16894043898758027805, name: "/device:GPU:0"
device_type: "GPU" memory_limit: 10088284160
locality {bus_id: 1 links {}}
incarnation: 17925533084010082620
physical_device_desc: "device: 0, name: GeForce RTX 3060, pci bus id: 0000:17:00.0, compute 
capability: 8.6"]

软件环境 #

使用Jupyter Notebook，已按照TensorFlow官方指引安装TensorFlow 2.6.0、CUDA和cuDNN。

问题现象 #

同时运行2个内核时，第一个内核可以正常运行Keras的Sequential模型，第二个内核运行相同代码时出现如下报错：

Attempting to perform BLAS operation using StreamExecutor without BLAS support
[[node sequential_3/dense_21/MatMul (defined at \AppData\Local\Temp/ipykernel_14764/3692363323.py:1) ]] [Op:__inference_train_function_7682]
函数调用栈：train_function

需求 #

不熟悉TensorFlow 1.x版本，询问单GPU下如何正常运行多个内核、共享GPU资源。

该报错的核心原因是Keras调用GPU运行时默认占用几乎全部显存，只需为每个Notebook单独配置显存上限即可，可按需调整如下代码中的memory_limit数值：

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=5120)])
  except RuntimeError as e:
    print(e)

内容的提问来源于stack exchange，提问作者MCPMH