RuntimeError: CUDA error: initialization

发布于:2025-02-11 ⋅ 阅读:(45) ⋅ 点赞:(0)

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)


网站公告

今日签到

点亮在社区的每一天
去签到