python如何指定多个gpu

Python指定多个GPU的方式主要有：使用CUDA_VISIBLE_DEVICES环境变量、使用TensorFlow的策略、使用PyTorch的DataParallel、使用Horovod。
其中，最常用的一种方法是使用环境变量CUDA_VISIBLE_DEVICES。通过设置该环境变量，可以控制Python程序使用哪些GPU。

具体方式如下：

在程序开始部分，设置环境变量，例如：

import os  
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

这行代码的意思是将GPU 0和1分配给当前运行的Python程序。接下来，我们会详细介绍其他几种方法。

一、CUDA_VISIBLE_DEVICES环境变量

设置环境变量CUDA_VISIBLE_DEVICES是最简单且最常用的方法之一，它适用于大多数深度学习框架。通过设置该环境变量，您可以控制程序使用哪些GPU。

设置环境变量

在Linux下，可以通过命令行设置环境变量：

export CUDA_VISIBLE_DEVICES=0,1

在Windows下：

set CUDA_VISIBLE_DEVICES=0,1

在Python程序中，可以通过os.environ设置环境变量：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

通过这种方式，您可以指定程序使用的GPU。

二、TensorFlow的策略

TensorFlow提供了多种策略来管理和使用多个GPU。最常用的策略是tf.distribute.MirroredStrategy，它可以自动将您的模型和数据分布到多个GPU上。

使用`tf.distribute.MirroredStrategy`

以下是一个简单的示例，展示如何使用tf.distribute.MirroredStrategy：

import tensorflow as tf
创建一个MirroredStrategy
strategy = tf.distribute.MirroredStrategy()
在策略范围内定义模型和训练步骤
with strategy.scope():
    # 定义模型
    model = tf.keras.models.Sequential([
        tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    # 编译模型
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    # 训练模型
    model.fit(train_dataset, epochs=5)

在这个示例中，tf.distribute.MirroredStrategy会自动将模型复制到所有可用的GPU上，并在每个设备上运行训练步骤。

三、PyTorch的DataParallel

PyTorch提供了torch.nn.DataParallel模块，用于在多个GPU上并行地训练模型。该模块会将输入数据分割到多个GPU上，并在每个设备上运行模型的一个副本。

使用`torch.nn.DataParallel`

以下是一个简单的示例，展示如何使用torch.nn.DataParallel：

import torch
import torch.nn as nn
import torch.optim as optim
定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(784, 10)
    def forward(self, x):
        return self.fc(x)
model = SimpleModel()
将模型包装到DataParallel中
model = nn.DataParallel(model)
将模型移动到GPU
model = model.cuda()
定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
训练模型
for data, target in train_loader:
    data, target = data.cuda(), target.cuda()
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

在这个示例中，nn.DataParallel会自动将输入数据分割到多个GPU上，并在每个设备上运行模型的一个副本。

四、Horovod

Horovod是一个用于分布式深度学习的开源库，它可以帮助您在多个GPU和多台机器上高效地训练模型。Horovod支持TensorFlow、Keras、PyTorch等多个深度学习框架。

安装Horovod

首先，您需要安装Horovod：

pip install horovod

使用Horovod与TensorFlow

以下是一个简单的示例，展示如何使用Horovod与TensorFlow：

import tensorflow as tf
import horovod.tensorflow.keras as hvd
初始化Horovod
hvd.init()
将TensorFlow的所有GPU可见性限制为当前进程
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
构建模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dense(10, activation='softmax')
])
编译模型
optimizer = tf.keras.optimizers.Adam()
optimizer = hvd.DistributedOptimizer(optimizer)
model.compile(loss='sparse_categorical_crossentropy',
              optimizer=optimizer,
              metrics=['accuracy'])
训练模型
model.fit(train_dataset, epochs=5, callbacks=[hvd.callbacks.BroadcastGlobalVariablesCallback(0)])

在这个示例中，Horovod会自动处理数据并行和通信问题，使得您可以在多个GPU上高效地训练模型。

五、其他常用方法

除了上述几种方法，还有一些其他常用的方法来在Python中使用多个GPU。例如，您可以使用NVIDIA的NCCL库来进行多GPU通信，或者使用Dask库来进行分布式计算。

使用NCCL

NCCL（NVIDIA Collective Communication Library）是一个高效的多GPU通信库，专为深度学习和大数据计算设计。它提供了高效的数据传输和通信原语，使得在多个GPU之间进行数据交换变得更加高效。

以下是一个简单的示例，展示如何使用NCCL：

import torch
import torch.distributed as dist
初始化NCCL
dist.init_process_group(backend='nccl')
获取当前进程的rank和world_size
rank = dist.get_rank()
world_size = dist.get_world_size()
创建一个张量
tensor = torch.ones(10).cuda()
进行AllReduce操作
dist.all_reduce(tensor)
输出结果
print('Rank:', rank, ' Tensor:', tensor)

在这个示例中，NCCL会自动处理多GPU之间的数据传输，使得AllReduce操作可以高效地在多个GPU上执行。

使用Dask

Dask是一个用于并行计算的Python库，它可以帮助您在多核CPU和多GPU上高效地执行计算任务。Dask提供了高层次的API，使得您可以轻松地进行分布式计算。

以下是一个简单的示例，展示如何使用Dask：

import dask.array as da
import dask.distributed as dd
创建Dask客户端
client = dd.Client()
创建一个Dask数组
x = da.random.random((10000, 10000), chunks=(1000, 1000))
进行计算
result = x.mean().compute()
输出结果
print('Mean:', result)

在这个示例中，Dask会自动将计算任务分配到多个CPU核和GPU上，使得计算任务可以高效地执行。

六、总结

在本文中，我们介绍了在Python中使用多个GPU的几种常用方法，包括使用CUDA_VISIBLE_DEVICES环境变量、使用TensorFlow的策略、使用PyTorch的DataParallel、使用Horovod以及其他常用方法。通过这些方法，您可以充分利用多个GPU的计算能力，加速深度学习模型的训练和推理过程。

无论您选择哪种方法，都需要根据具体的应用场景和需求进行调整和优化。希望本文能够帮助您更好地理解和使用多GPU技术，提高深度学习模型的训练效率和性能。