将深度学习代码同时运行在多个GPU显卡上

将深度学习代码同时运行在多个GPU显卡上

步骤1

设置os.environ[‘CUDA_VISIBLE_DEVICES’] 指定GPU

# 将import torch放在os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'后面
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
import torch
print(torch.cuda.device_count()) #2说明指令生效

# 不能将import torch放在前面

步骤2(任选其一)

1.使用DataParallel(适合PyTorch)

import torch
import torch.nn as nn

# 假设有个model
model = Yourmodel()

# 将模型放到多个GPU上(例如0，1)
model = nn.DataParallel(model,device_ids=[0,1])
# 将模型移到GPU上
model = model.cuda()
# 向前传播会自动在多个GPU上并行处理
output = model(input_data)

2.使用DistributedDataParallel(适合分布式训练，支持多机多卡)

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

def train(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    
    # 设置 GPU 设备
    torch.cuda.set_device(rank)
    
    # 假设你有一个模型
    model = YourModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    # 模型训练等代码
    output = ddp_model(input_data)

# 在两个 GPU 上启动进程
world_size = 2
mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

3. TensorFlow: 手动指定 GPU 并行

   import tensorflow as tf
   
   # 创建分布式策略
   strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
   
   with strategy.scope():
       model = YourModel()
       model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
   
   # 训练代码会自动并行运行在多个 GPU 上
   model.fit(train_data, epochs=5)

4. 使用 CUDA 的多 GPU 手动操作

   对于更灵活的多 GPU 控制，你可以手动将模型拆分并将不同部分放在不同的 GPU 上。不过这种方法复杂度较高，适用于需要精细控制 GPU 使用的情况。

注意事项

确保每个 GPU 的内存足够，特别是分布式训练时。
对于 DistributedDataParallel 和 MirroredStrategy，可能需要调整 batch size 以平衡每个 GPU 的负载。