以下是几种指定当前Python程序使用多张显卡的方法： #### PyTorch中指定多张显卡 在PyTorch里，可借助`DataParallel`或者`DistributedDataParallel`来实现多显卡训练。 ##### 使用`DataParallel` `DataParallel`能自动把模型复制到所有可用的GPU上，并且在每个GPU上并行处理数据。 ```python import torch import torch.nn as nn # 定义模型 model = YourModel() # 检查是否有可用的GPU if torch.cuda.device_count() > 1: print(f"Using {torch.cuda.device_count()} GPUs") model = nn.DataParallel(model) # 将模型移动到GPU model.cuda() ``` ##### 使用`DistributedDataParallel` `DistributedDataParallel`通常比`DataParallel`更高效，特别是在使用多个节点或者多张显卡时。 ```python import torch import torch.distributed as dist import torch.multiprocessing as mp import torch.nn as nn import torch.optim as optim from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '12355' # initialize the process group dist.init_process_group("nccl", rank=rank, world_size=world_size) def cleanup(): dist.destroy_process_group() def run(rank, world_size): setup(rank, world_size) # create model and move it to GPU with id rank model = YourModel().to(rank) ddp_model = DDP(model, device_ids=[rank]) # define loss function and optimizer loss_fn = nn.MSELoss() optimizer = optim.SGD(ddp_model.parameters(), lr=0.001) # training loop for epoch in range(num_epochs): # training code here pass cleanup() if __name__ == "__main__": world_size = torch.cuda.device_count() mp.spawn(run, args=(world_size,), nprocs=world_size, join=True) ``` #### TensorFlow中指定多张显卡 在TensorFlow里，能够使用`tf.distribute.MirroredStrategy`来指定使用多张显卡。 ```python import tensorflow as tf # 创建一个MirroredStrategy对象 strategy = tf.distribute.MirroredStrategy() print('Number of devices: {}'.format(strategy.num_replicas_in_sync)) # 在策略范围内定义模型 with strategy.scope(): model = tf.keras.Sequential([...]) model.compile(...) # 训练模型 model.fit(...) ``` #### 通过环境变量指定多张显卡 可以通过设置`CUDA_VISIBLE_DEVICES`环境变量来指定程序可见的显卡，这样程序就会使用这些显卡。 ```python import os os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" # 指定使用显卡0, 1, 2, 3 ```