在 PyTorch 中强制 GPU 内存限制

问题描述

有没有办法强制为特定 Pytorch 实例可用的 GPU 内存量设置最大值?例如，我的 GPU 可能有 12Gb 可用，但我想为特定进程分配最大 4Gb.

Is there a way to force a maximum value for the amount of GPU memory that I want to be available for a particular Pytorch instance? For example, my GPU may have 12Gb available, but I'd like to assign 4Gb max to a particular process.

推荐答案

更新 (04-MAR-2021):现在是 在稳定的 1.8.0 版本的 PyTorch 中可用.此外，在文档中

Update (04-MAR-2021): it is now available in the stable 1.8.0 version of PyTorch. Also, in the docs

原始答案如下.

此功能请求已合并到 PyTorch master 分支.然而，没有在稳定版本中引入.

This feature request has been merged into PyTorch master branch. Yet, not introduced in the stable release.

作为memory_set_引入代码>

为进程设置内存分数.该分数用于将缓存分配器限制为在 CUDA 设备上分配的内存.允许的值等于总的可见内存乘数.如果尝试在进程中分配超过允许的值，将引发 out of分配器内存错误.

您可以查看 as">测试

You can check the tests as usage examples.

这篇关于在 PyTorch 中强制 GPU 内存限制的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！