当前位置：首页 > article >正文

RuntimeError: CUDA error: initialization

article 2025/1/4 6:11:07

RuntimeError: CUDA error: initialization

cuda初始化出问题了，这是因为在python多线程跑gpu代码程序时先对cuda进行操作，然后在跑gpu代码时就没有cuda可用了。
在main的主程序代码加一行代码就可以了，用来获取cuda，在代码中只能使用一次：

import multiprocessing as mp

if __name__ == "__main__":
    mp.set_start_method('spawn')

多进程推理代码：

import os

os.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=1234
    generator = 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)

查看全文

http://www.kler.cn/a/459951.html