一文读懂 PyTorch 显存管理机制

文章预览

作者丨米阿罗@知乎（已授权转载） 来源丨https://zhuanlan.zhihu.com/p/486360176 编辑丨小书童,集智书童 1、背景介绍 剖析 PyTorch 显存管理机制主要是为了减少 「显存碎片化」 带来的影响。一个简单示例为： 如上图所示，假设当前想分配 800MB 显存，虽然空闲的总显存有 1000MB，但是上方图的空闲显存由地址不连续的两个 500MB 的块组成，不够分配这 800MB 显存；而下方的图中，如果两个 500MB 的空闲块地址连续，就可以通过显存碎片的整理组成一个 1000MB 的整块，足够分配 800MB。上方图的这种情况就被称为 「显存碎片化」 。 「解决方法」 ：当有多个 Tensor 可以被释放时，可以优先释放那些在内存空间上前后有空闲内存的 Tensor。这样便于 PyTorch 整理这些空闲块组成的碎片，让三个块组成一整个更大的块。 「核心问题/前提」 ：能否以非常小的代价（ O(1) 或 O(lo ………………………………