Python指定GPU|极客教程

Python指定GPU

在使用Python进行深度学习或机器学习任务时，通常会涉及到使用GPU来加速模型的训练过程。在一台拥有多个GPU的机器上，我们可以通过指定GPU来为特定任务分配计算资源，以实现更高效的训练过程。本文将介绍如何在Python中指定GPU来进行深度学习任务以及一些常见的实例代码。

检查可用的GPU

在使用Python进行深度学习任务之前，首先需要检查机器上是否安装了合适的GPU驱动和CUDA。可以使用以下代码来检查系统中可用的GPU设备：

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        print("GPU device name:", gpu.name)
else:
    print("No GPU devices found.")

运行以上代码，输出将显示系统中可用的GPU设备名称。如果系统中没有GPU设备，将显示”No GPU devices found.”。

指定GPU设备

在训练模型时，有时候我们希望指定特定的GPU设备进行训练。可以使用以下代码来指定GPU设备：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 指定使用第一个GPU设备

运行以上代码，将只有第一个GPU设备可用进行训练。如果要指定使用多个GPU设备，可以将设备索引用逗号分隔开。

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"  # 指定使用第一个和第二个GPU设备

示例代码

下面我们通过一个简单的示例代码来演示如何在Python中指定GPU来进行深度学习任务。假设我们要使用TensorFlow来训练一个简单的神经网络模型，以下是示例代码：

import tensorflow as tf
from tensorflow.keras import layers

# 指定使用第一个GPU设备
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# 构建一个简单的神经网络模型
model = tf.keras.Sequential([
    layers.Dense(64, activation='relu', input_shape=(100,)),
    layers.Dense(64, activation='relu'),
    layers.Dense(10)
])

# 编译模型
model.compile(optimizer=tf.keras.optimizers.Adam(0.001),
              loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=10, batch_size=32)

在以上示例代码中，我们指定了只使用第一个GPU设备来训练神经网络模型。你也可以根据自己的需求来指定训练过程中使用的GPU设备。