【Pytorch】在多进程中使用 CUDA

Python 原生自带的多进程库不支持在子进程中调用 CUDA 进行加速运算。因此,本文介绍了使用 Pytorch 中的 multiprocessing 库实现在子进程中调用 CUDA 的方法。

在使用 Python 原生自带的 multiprocessing 库中,若在子进程中调用 CUDA,则会报错:

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

如果根据提示信息,直接使用

torch.multiprocessing.set_start_method("spawn")

则会报错:

RuntimeError(‘context has already been set’)

这是因为想要实现在多进程中调用 CUDA,需要先新建一个 context 对象,并在实际使用中用这个 context 对象替换 torch.multiprocessing,例如:

context = torch.multiprocessing.get_context("spawn")
...
process = context.Process(target=..., args=...)

此时,在子进程中就可以放心地调用 CUDA 了。

相关推荐

  1. Pytorch进程使用 CUDA

    2024-01-11 07:46:05       37 阅读
  2. PyTorch常用的工具(5)使用GPU加速:CUDA

    2024-01-11 07:46:05       38 阅读
  3. pytorch 支持更的 onnx 算子

    2024-01-11 07:46:05       48 阅读
  4. 关于PyTorch使用cudnn.benchmark= True

    2024-01-11 07:46:05       39 阅读

最近更新

  1. TCP协议是安全的吗?

    2024-01-11 07:46:05       16 阅读
  2. 阿里云服务器执行yum,一直下载docker-ce-stable失败

    2024-01-11 07:46:05       16 阅读
  3. 【Python教程】压缩PDF文件大小

    2024-01-11 07:46:05       15 阅读
  4. 通过文章id递归查询所有评论(xml)

    2024-01-11 07:46:05       18 阅读

热门阅读

  1. 数据库基础5

    2024-01-11 07:46:05       31 阅读
  2. 5个Linux归档命令

    2024-01-11 07:46:05       34 阅读
  3. 初探Spark SQL catalog缓存机制

    2024-01-11 07:46:05       35 阅读
  4. K8S容器编排高级应用

    2024-01-11 07:46:05       35 阅读
  5. 微服务治理:什么是微服务生命周期管理 (MLM)?

    2024-01-11 07:46:05       30 阅读
  6. MySQL自动备份脚本

    2024-01-11 07:46:05       39 阅读