TensorFlow GPU 配置教程

本文将为您介绍如何在您的系统上配置 TensorFlow 以支持 GPU 加速。TensorFlow 是一个强大的开源机器学习库，通过 GPU 加速可以显著提高训练和推理的速度。

系统要求

• 操作系统：Linux、macOS 或 Windows
• CUDA：至少 CUDA 10.0
• cuDNN：与 CUDA 版本兼容的 cuDNN 版本
• Python：Python 3.5 或更高版本

安装步骤

1. 安装 CUDA 和 cuDNN：请访问 NVIDIA 官网下载并安装与您的 GPU 兼容的 CUDA 和 cuDNN。

2. 安装 Python：确保您的系统已安装 Python 3.5 或更高版本。

3. 安装 TensorFlow：使用以下命令安装 TensorFlow：

pip install tensorflow-gpu

配置 GPU

1. 设置环境变量：在您的系统环境中设置以下变量：

export CUDA_HOME=/usr/local/cuda
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
export PATH=$CUDA_HOME/bin:$PATH

2. 验证安装：在 Python 中运行以下代码以验证 TensorFlow 是否已正确配置：

import tensorflow as tf

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

示例代码

以下是一个简单的 TensorFlow GPU 加速示例：

import tensorflow as tf

# 创建一个简单的模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(32,)),
    tf.keras.layers.Dense(1)
])

# 编译模型
model.compile(optimizer='adam',
              loss='mean_squared_error')

# 训练模型
model.fit(tf.random.normal([1000, 32]), tf.random.normal([1000, 1]), epochs=10)

扩展阅读

如果您想了解更多关于 TensorFlow 的信息，可以访问我们的 TensorFlow 教程页面。

[center]