每次安装TensorFlow最大的痛就是跟CUDA的兼容问题。费半天劲安好了Nvidia-driver、 CUDA 、CuDNN、TensorRT,再安装TensorFlow,然后发现TensorFlow还是不能启用GPU。
今天发现TensorFlow已经发布2.15.0版了,更新介绍说这次可以通过pip一次把TensorFlow和相应的CUDA都安装好了。试了一下,果然省心多了。
官方安装指导:Install TensorFlow with pip (google.cn)
注意一定要在页面右上角把语言切换成英文,中文的内容(截止2023.12.04日)还没更新。
主要的安装步骤是:
python3 -m pip install --extra-index-url https://pypi.nvidia.com tensorrt-bindings==8.6.1 tensorrt-libs==8.6.1
python3 -m pip install -U tensorflow[and-cuda]
安装过程会下载和安装很多个组件。我这边下载速度大概800k/s,安装用了一个小时。安装完成后用下面的语句验证:
python3 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"