内存与性能的联系

PyTorch的内存机制如下:

Tensor对象管理：PyTorch在单算子场景下，算子的下发执行是动态的（dynamic eager execution），同时Tensor的创建和销毁也是动态的，因此执行的某一时刻并不知道当前已经创建的Tensor何时使用和释放。在eager的场景下，PyTorch依赖Python垃圾回收机制，通过引用计数的方式跟踪管理每个Tensor对象，引用计数达到0时销毁对象释放内存，该机制确保Tensor不需要时准确释放内存。

CachingAllocator：在网络的计算过程中，需要动态的为每个算子申请输出Tensor内存，Tensor的生命周期结束后释放内存，存在大量的内存申请释放。为提高内存申请释放速度，PyTorch在Allocator中通过内存池缓存内存块，减少内存申请释放API的调用。内存申请时先从内存池中寻找可用的内存块，内存池中找不到调用API（cudaMalloc/aclrtMallocAlign32）向驱动申请；内存释放时归还给内存池，并不调用API（cudaFree/aclrtFree）真正释放归还内存给驱动。内存池中持有的内存通常不还给驱动，只有在内存申请失败或者用户主动调用API（torch.npu.empty_cache()）时释放回驱动。

总的来说，PyTorch在内存上基于eager模式实现，eager模式的特点是如果没有必要，不会主动释放已经申请的内存，同时，PyTorch默认模式下是整块申请内存，因此，很容易出现碎片化的内存，而当模型当前内存不足够时，会触发内存重整操作，即将碎片化内存（包括之前申请但并未被表示的内存块）整理，得到大块连续内存。如果重整得到的内存能够满足模型训练需要，那么训练就会进行下去，但是这种内存重整操作需要花费大量时间，频繁触发内存重整操作会对模型性能产生很大影响。因此，本文就内存问题提出一些针对化建议。