如何处理训练过程中出现OOM（显存不足）

【现象描述】

GPU上网络运行过程中出现OOM（显存不足）

【原因分析】

出现该现象，可能原因如下：

1.训练batchsize过大。

2.输入数据的shape是变化的（输入数据动态shape）。

3.输出结果的tensor保存起来了。

4.网络中的算子里出现显存泄漏（算子里每次launch都申请显存，并且不释放）。

【解决方法】

步骤1：排查训练的batchsize是否过大，可以逐步缩小batch，如果觉得batch不够大，还是出现OOM，则排除这个原因。

步骤2：确认输入数据是否是动态shape，当前动态shape还不支持输入shape可变，因为针对这种场景，每次step训练都会编译新图，导致显存不断申请最后OOM，可以通过图的个数是不是不断増长来判断。

步骤3：输出tensor里会挂接device上的地址，方便print时同步device数据输出，在tensor析构的时候会释放device地址，因此如果每次step训练都把输出tensor保存起来的话，导致tensor不会析构，随着训练step的增加则会OOM。

步骤4：前面3个步骤确认都没有问题的话，则有可能是算子实现bug，可以排查下是否有新增算子，算子里是否有申请显存的操作导致显存泄漏。

案例代码：

```

bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
              const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
    T *input = GetDeviceAddress<T>(inputs, 0);
    S *indices = GetDeviceAddress<S>(inputs, 1);
    T *updates = GetDeviceAddress<T>(inputs, 2);
    T *output = GetDeviceAddress<T>(outputs, 0);

    const size_t indices_len = sizeof(S) * out_strides_.size();
    void *indices_stride_work = device::gpu::GPUMemoryAllocator::GetInstance().AllocTensorMem(indices_len);
    if (indices_stride_work == nullptr) {
      MS_LOG(EXCEPTION) << "Failed to alloc indices_stride_work, size: " << indices_len;
    }

```

---》每次launch都调用内存池接口AllocTensorMem申请内存，导致随着训练step增加出现OOM

【建议与总结】

1.输出tensor一般不建议保存起来，如果需要保存，可以保存asnumpy后的对象，tensor经过asnumpy后拷贝device到host后会释放掉device的地址。

2.算子实现逻辑不应该出现申请显存的逻辑，一般都可以通过workspace来实现临时device地址存储的目的。

【相关参考文档】

无

最后

以上就是结实小蘑菇为你收集整理的如何处理训练过程中出现OOM（显存不足）的全部内容，希望文章能够帮你解决如何处理训练过程中出现OOM（显存不足）所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。