TensorFlow 2.1启用GPU加速

2020/4/30 by Qin Zhiqiang in 深度学习 TensorFlow, GPU, CUDA, AI

Python、TensorFlow、NVIDIA GPU驱动、CUDA工具包等软件的版本需要相互匹配，可参考官方的指导“Install TensorFlow 2”与"软件要求"。TensorFlow 2.x已不需独立的TensorFlow-GPU版本。

软件匹配版本：

Python 3.7
NVIDIA GPU驱动 445.87
CUDA工具包 10.1
cuDNN SDK 7.6
TensorFlow 2.1

软件安装

Python: https://www.python.org/ftp/python/3.7.7/python-3.7.7-amd64.exe
NVIDIA GPU驱动: https://www.nvidia.com/drivers
CUDA工具包: https://developer.nvidia.com/cuda-toolkit-archive，安装指导见 https://docs.nvidia.com/cuda/cuda-installation-guide-microsoft-windows/index.html。
cuDNN SDK: https://developer.nvidia.com/cudnn；需解压文件到CUDA的指定目录，见 https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#install-windows
TensorFlow 2.1: 国内pip安装比较慢，为了避免网络中断，先安装部分依赖软件包。看到下载比较慢时，可以使用迅雷下载如numpy、h5py、scipy、tensorflow 2.1等软件包，再本地安装，如：pip install .\numpy-1.18.3-cp37-cp37m-win_amd64.whl。
```
python -m pip install --upgrade pip
python -m pip install wheel

pip install numpy
pip install pandas
pip install scipy

pip install protobuf grpcio h5py tensorflow_estimator

pip install tensorflow==2.1
```
whl文件地址：numpy、pandas、scipy、protobuf、grpcio、h5py、tensorflow_estimator、tensorflow。

运行MNIST样例

import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import tensorflow as tf

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test,  y_test, verbose=2)

TensorFlow第一次启用CPU加速时，需要等待很久。同时，mnist.load_data()首次下载数据也耗时。

修正CUDNN_STATUS_INTERNAL_ERROR

运行较大的模型时，可能出现CUDNN_STATUS_INTERNAL_ERROR。在脚本开始处添加如下代码，可避免该错误。

# Fix CUDNN_STATUS_INTERNAL_ERROR
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession

config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)

禁用GPU加速

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"