RuntimeError: CUDA error: initialization
cuda初始化出问题了,这是因为在python多线程跑gpu代码程序时先对cuda进行操作,然后在跑gpu代码时就没有cuda可用了。
在main的主程序代码加一行代码就可以了,用来获取cuda,在代码中只能使用一次:
import multiprocessing as mpif __name__ == "__main__":mp.set_start_method('spawn')
多进程推理代码:
import osos.environ['CUDA_VISIBLE_DEVICES']='0'import torch
import multiprocessing# 定义每个进程要执行的函数,这里简单做一个张量求和计算示例
def process_task(gpu_id, tensor_data):# 设置当前进程可见的CUDA设备# os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)print("gpu_id",gpu_id)device= torch.device(f"cuda:{gpu_id}")seed=1234generator = torch.Generator(device).manual_seed(seed)# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")tensor = tensor_data.to(device)result = tensor.sum()return result.item()if __name__ == "__main__":num_processes = 5 # 定义要启动的进程数量,这里设置为2,可根据实际GPU数量等情况调整gpu_ids = [2,2,5] # 对应每个进程使用的GPU设备编号,需根据实际系统中的GPU情况安排tensor_list = [torch.randn(5, 5) for _ in range(num_processes)] # 模拟每个进程要处理的张量数据with multiprocessing.Pool(num_processes) as pool:args_list = [(gpu_id, tensor) for gpu_id, tensor in zip(gpu_ids, tensor_list)]results = pool.starmap(process_task, args_list)print("各个进程的计算结果:", results)