在Google Colab中安装CUDA与cuDNN及解决训练脚本报错问题

我在Google Colab上训练目标检测数据集，已经把数据集上传到Google Drive并在Colab中成功调用，但运行以下训练脚本时出现错误：

!python3 /content/drive/tensorflow1/models/research/object_detection/train.py --logtostderr --train_dir=/content/drive/tensorflow1/models/research/object_detection/training/ --pipeline_config_path=/content/drive/tensorflow1/models/research/object_detection/training/faster_rcnn_inception_v2_pets.config

报错信息如下：

Traceback (most recent call last): File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/pywrap_tensorflow.py", line 58, in <module> from tensorflow.python.pywrap_tensorflow_internal import * File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/pywrap_tensorflow_internal.py", line 28, in <module> _pywrap_tensorflow_internal = swig_import_helper() File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/pywrap_tensorflow_internal.py", line 24, in swig_import_helper _mod = imp.load_module('_pywrap_tensorflow_internal', fp, pathname, description) File "/usr/lib/python3.6/imp.py", line 243, in load_module return load_dynamic(name, filename, file) File "/usr/lib/python3.6/imp.py", line 343, in load_dynamic return _load(spec) ImportError: libcublas.so.9.0: cannot open shared object file: No such file or directory During handling of the above exception, another exception occurred: Traceback (most recent call last): File "/content/drive/tensorflow1/models/research/object_detection/train.py", line 47, in <module> import tensorflow as tf File "/usr/local/lib/python3.6/dist-packages/tensorflow/__init__.py", line 24, in <module> from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/__init__.py", line 49, in <module> from tensorflow.python import pywrap_tensorflow File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/pywrap_tensorflow.py", line 74, in <module> raise ImportError(msg) ImportError: Traceback (most recent call last): File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/pywrap_tensorflow.py", line 58, in <module> from tensorflow.python.pywrap_tensorflow_internal import * File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/pywrap_tensorflow_internal.py", line 28, in <module> _pywrap_tensorflow_internal = swig_import_helper() File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/pywrap_tensorflow_internal.py", line 24, in swig_import_helper _mod = imp.load_module('_pywrap_tensorflow_internal', fp, pathname, description) File "/usr/lib/python3.6/imp.py", line 243, in load_module return load_dynamic(name, filename, file) File "/usr/lib/python3.6/imp.py", line 343, in load_dynamic return _load(spec) ImportError: libcublas.so.9.0: cannot open shared object file: No such file or directory Failed to load the native TensorFlow runtime. See https://www.tensorflow.org/install/install_sources#common_installation_problems for some common reasons and solutions. Include the entire stack trace above this error message when asking for help.

请问是否需要先将CUDA9或cuDNN安装/上传至Google Drive才能在Colab中解决该问题？如何修复这些错误？

别担心，你完全不需要自己把CUDA或者cuDNN上传到Google Drive来解决这个问题。这个错误的核心原因是你当前使用的TensorFlow 1.x版本依赖CUDA 9.0，但Colab默认的GPU环境已经升级到了更高版本的CUDA（比如10.x或11.x），两者不兼容导致的。下面是具体的修复步骤：

1. 确认Colab当前的CUDA版本 #

首先运行这条命令查看Colab分配给你的GPU对应的CUDA版本：

!nvidia-smi

你会看到类似这样的输出，重点看右上角的CUDA Version字段，比如可能是10.1或者11.8。

2. 安装与CUDA版本兼容的TensorFlow 1.x版本 #

根据CUDA版本选择对应的TensorFlow版本：

• 如果CUDA版本是10.0/10.1：安装TensorFlow 1.15（这是TF1.x的最后一个稳定版本，兼容CUDA 10.x）
• 如果CUDA版本是9.0：可以安装TensorFlow 1.13（不过现在Colab很少分配CUDA9.0的环境了）

运行以下命令卸载当前的TensorFlow并安装兼容版本：

!pip uninstall tensorflow -y
!pip install tensorflow==1.15

3. 确保Object Detection API的环境路径配置正确 #

在运行训练脚本前，需要确保Object Detection API的相关路径被添加到Python的环境变量中，否则会出现模块找不到的问题。运行以下代码：

import sys
sys.path.append('/content/drive/tensorflow1/models/research')
sys.path.append('/content/drive/tensorflow1/models/research/slim')

另外，还要确保你已经编译了API的proto文件（如果之前没做过的话）：

%cd /content/drive/tensorflow1/models/research
!protoc object_detection/protos/*.proto --python_out=.

4. 重新运行训练脚本 #

完成以上步骤后，再次执行你原来的训练命令，应该就能正常运行了：

!python3 /content/drive/tensorflow1/models/research/object_detection/train.py --logtostderr --train_dir=/content/drive/tensorflow1/models/research/object_detection/training/ --pipeline_config_path=/content/drive/tensorflow1/models/research/object_detection/training/faster_rcnn_inception_v2_pets.config

额外建议 #

如果你打算长期使用目标检测训练，也可以考虑迁移到TensorFlow 2.x版本的Object Detection API，它对Colab的最新环境兼容性更好，而且性能和易用性都有提升。不过这需要你调整pipeline配置文件和训练脚本，适合有一定基础后再尝试。

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