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

Linux下用多进程在GPU上跑Pytorch模型问题

先看一段代码

import concurrent.futures
import torch

device = "cuda"

model = torch.nn.Linear(20, 30)
model.to(device)


def exec(v):
    input = torch.randn(128, 20).to(device)
    output = model(input)
    return v


if __name__ == '__main__':
    with concurrent.futures.ProcessPoolExecutor(max_workers=5) as executor:
        for x in executor.map(exec, range(10)):
            print(x)

这段代码尝试在多个进程下并行运行一个pytorch网络。在Linux下运行会遇到如下错误:

RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

还有一种可能的情况是死锁,这可能和模型大小有关。总之,上面这段代码是无法正常于行的。

问题原因与python multiprocess的启动方式有关。Python的multiprocess一共有三种启动方式:spawnforkforkserver。而CUDA runtime是不支持fork的。不幸的是,除了macOS以外的POSIX系统的默认方式都是fork。

知道了原因,解决方法就非常简单了,那就是显式设置multiprocess的启动方式

import concurrent.futures
import torch
import multiprocessing as mp

device = "cuda"

model = torch.nn.Linear(20, 30)
model.to(device)


def exec(v):
    input = torch.randn(128, 20).to(device)
    output = model(input)
    return v


if __name__ == '__main__':
    mp.set_start_method('spawn')
    with concurrent.futures.ProcessPoolExecutor(max_workers=5) as executor:
        for x in executor.map(exec, range(10)):
            print(x)

 这次就能正常运行了

0
1
2
3
4
5
6
7
8
9

Reference:

Multiprocessing best practices — PyTorch 2.6 documentation

multiprocessing — Process-based parallelism — Python 3.13.2 documentation


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

相关文章:

  • python -面试题--算法
  • 安科瑞ACCU-100微电网协调控制器:助力绿色能源系统运行
  • JVM之基础知识
  • 以实现生产制造、科技研发、人居生活等一种或多种复合功能的智慧油站开源了
  • 蓝桥杯 互质数的个数
  • Axure RP下载安装和简单使用教程
  • 浙江大学:DeepSeek行业应用案例集(153页)(文末可下载PDF)
  • Python爬虫:从人民网提取视频链接的完整指南
  • 使用1Panel一键搭建WordPress网站的详细教程(全)
  • 力扣hot100二刷——链表
  • 【Linux 指北】常用 Linux 指令汇总
  • 强化学习(赵世钰版)-学习笔记(7.时序差分学习)
  • Centos离线安装openssl
  • DeepSeek-prompt指令-当DeepSeek答非所问,应该如何准确的表达我们的诉求?
  • 单体架构、微服务组件与解决方案
  • 【计量地理学】实验二 经典统计分析方法
  • ⭐算法OJ⭐汉明距离【位操作】(C++ 实现)Hamming Distance
  • Ubuntu24.04 LTS 版本 Linux 系统在线和离线安装 Docker 和 Docker compose
  • Spring MVC面试题(一)
  • Linux Shell 脚本编程极简入门指南