### Python中使用多张GPU进行并行计算的方法 在Python中，可以借助多种库来实现多张GPU的并行计算。以下是几种常见的方法及其具体实现方式： #### 1. **使用PyTorch进行多GPU并行计算** PyTorch提供了`torch.nn.DataParallel`模块和分布式训练工具`torch.distributed`，用于简化多GPU环境下的模型训练。 - `DataParallel`是一种简单的解决方案，适用于小型项目或快速原型设计。它会自动将输入数据分发到不同的GPU上，并收集结果[^4]。 ```python import torch import torch.nn as nn class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.fc = nn.Linear(100, 1) def forward(self, x): return self.fc(x) model = SimpleModel() if torch.cuda.device_count() > 1: model = nn.DataParallel(model) # 自动分配到多个GPU model.to('cuda') ``` - 对于更复杂的场景，推荐使用`torch.distributed.launch`脚本启动多进程模式。这允许开发者手动控制每张GPU上的操作逻辑[^3]。 ```bash CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch \ --nproc_per_node=4 --use_env train_script.py ``` --- #### 2. **使用TensorFlow/Keras进行多GPU并行计算** TensorFlow支持通过策略（Strategy）API配置多GPU环境。其中最常用的是`tf.distribute.MirroredStrategy`，它可以轻松管理同步更新权重的过程[^1]。 ```python import tensorflow as tf strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"]) with strategy.scope(): model = tf.keras.Sequential([ tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(10, activation='softmax') ]) model.compile(optimizer=tf.keras.optimizers.Adam(), loss=tf.keras.losses.SparseCategoricalCrossentropy(), metrics=['accuracy']) ``` --- #### 3. **使用CuPy或其他专用库** 对于不需要深度学习框架的任务，可以直接利用CuPy等库完成数值运算的多GPU扩展。虽然这种方法灵活性较高，但需要更多手工编码以协调不同设备间的通信。 ```python import cupy as cp # 初始化两个数组分别位于不同的GPU上 with cp.cuda.Device(0): array_on_gpu0 = cp.random.rand(1000, 1000) with cp.cuda.Device(1): array_on_gpu1 = cp.random.rand(1000, 1000) # 执行跨GPU的操作 (需显式传输数据) result = array_on_gpu0.dot(cp.asnumpy(array_on_gpu1)) ``` --- #### 4. **基于pyCUDA的手工优化方案** 尽管不常见，某些特殊需求可能更适合采用低级别的pyCUDA接口直接编写CUDA核函数[^5]。这种方式性能最优但也最难维护。 ```python import pycuda.autoinit import pycuda.driver as drv from pycuda.compiler import SourceModule mod = SourceModule(""" __global__ void multiply_them(float *dest, float *a, float *b) { const int i = threadIdx.x; dest[i] = a[i] * b[i]; } """) multiply_them = mod.get_function("multiply_them") ... ``` 注意此路径仅适合熟悉底层硬件特性的专家级用户。 --- ### 总结 以上列举了几种主流技术路线供参考。实际应用时应综合考虑项目规模、团队技能水平以及目标平台特性等因素决定最佳实践方向。