解决Python中cuBLAS插件无法注册问题
在使用TensorFlow进行深度学习模型训练时,有时候会遇到"Unable to register cuBLAS factory Attempting to register factory for plugin cuBLAS"的错误提示,这个错误通常是由于cuBLAS插件未正确注册引起的。本篇文章将介绍如何通过代码来解决这个问题。
首先,需要在Python脚本中添加以下代码:
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
from tensorflow.python.platform import build_info
print(build_info.cuda_version_number)
运行这段代码后,如果控制台输出了CUDA的版本号,说明cuBLAS已经成功注册,可以继续进行模型训练了。如果还是出现上述错误,则需要按照以下步骤进行操作。
sudo apt-get remove --purge cuda
sudo apt-get remove --purge libcudnn*
sudo apt-get install cuda=10.0.130-1
sudo apt-get install libcudnn7-dev=7.6.3.30-1+cuda10.0
注意,需要根据自己的实际情况修改版本号。
pip ins