如果我想限制每个进程只能使用一张GPU，应该如何做？

如果你想要限制每个进程只使用一张GPU，通常在创建并行计算任务时，可以为每个进程分配特定的GPU。以下是一个使用Python和多进程（`multiprocessing`）的例子，它确保每个子进程运行时只会占用一个GPU：

import os
from torch.multiprocessing import Pool, Process
import torch

def setup_multiproc(rank, world_size):
    # 设置每个进程的CUDA设备
    os.environ['CUDA_VISIBLE_DEVICES'] = str(rank)
    torch.cuda.set_device(rank)

# 定义单个处理单元的函数，这里假设它需要GPU资源
def process_function(rank, data):
    setup_multiproc(rank, len(os.environ['CUDA_VISIBLE_DEVICES']))  # 确保每个进程单独设置GPU
    with torch.cuda.device(rank):
        # 在这里进行GPU相关的操作
        ...

# 创建一个进程池，其中world_size是你计划使用的GPU数量
world_size = torch.cuda.device_count()
pool = Pool(processes=world_size)

# 分配数据给每个进程
data_list = [data_for_process_i for i in range(world_size)]
results = pool.starmap(process_function, [(i, d) for i, d in enumerate(data_list)])

# 关闭进程池
pool.close()
pool.join()

在这个例子中，`setup_multiproc`函数负责设置每个进程的环境变量`CUDA_VISIBLE_DEVICES`，使其只看到一个GPU。然后每个进程在进入`process_function`时，会自动切换到正确的GPU设备上。