从句子转换器创建对象时 GPU 内存泄漏

Question

描述

我正在使用 Python 中的 sentence_transformers 库在 R 中创建一个嵌入句子的函数。

由于某些未知原因，在同一变量名下多次创建对象最终导致内存不足 space 无法分配转换器。重现：

sentence_transformers <- reticulate::import("sentence_transformers")
for (i in 1:10) {
  print(i)
  bert_encoder <- sentence_transformers$SentenceTransformer("bert-large-nli-stsb-mean-tokens")
}

但是，直接对 Python 执行相同的操作不会产生错误

from sentence_transformers import SentenceTransformer
for i in range(10):
    print(i)
    bert_encoder = SentenceTransformer("bert-large-nli-stsb-mean-tokens")
}

在 GPU 中分配的任何模型都会发生这种情况。在我的 NVIDIA GTX 1060 上它达到了第 4 个周期，但在较小的 GPU 上它崩溃得更早。一种临时解决方案是只在外部创建一次模型，然后将模型作为参数多次传递给函数，但我宁愿避免这种情况，因为它增加了一个额外的步骤，并且在任何情况下调用多个模型可能只是让它也崩溃。

预期行为

for 循环结束且没有错误

观察到的行为

py_call_impl(callable, dots$args, dots$keywords) 错误：运行时错误：CUDA 内存不足。尝试分配 20.00 MiB（GPU 0；6.00 GiB 总容量；2.95 GiB 已分配；16.11 MiB 可用；238.68 MiB 缓存）

尝试求解失败

提出的解决方案here
按照建议使用 numba
通过 reticulate::py_run_string() 在 Python 上显式声明变量，然后执行 del bert_encoder 并调用垃圾收集器

详情

Windows 10 家

Python 3.7.4

R 4.0.1

网纹 1.16

手电筒 1.3.1

张量流 2.2.0

变形金刚 2.11.0

sentence_transformers 0.2.6

Answer 1

好的，所以我将我的解决方案发布给遇到此问题的其他人。

每次调用模型后

sentence_transformers <- import("sentence_transformers")
encoder <- sentence_transformers$SentenceTransformer("bert-large-nli-stsb-mean-tokens")

我使用

释放GPU内存

  # Has this been done on a GPU?
  py <- reticulate::py_run_string("import torch
is_cuda_available = torch.cuda.is_available()")

  # Release GPU
  if (isTRUE(reticulate::py$is_cuda_available)) {

    tryCatch(reticulate::py_run_string("del encoder"),
             warning = function(e) {},
             error = function(e) {})

    tryCatch(rm(encoder),
             warning = function(e) {},
             error = function(e) {})

    gc(full = TRUE, verbose = FALSE)

    py <- reticulate::py_run_string("import torch
torch.cuda.empty_cache()")

  }

而且效果很好。

从句子转换器创建对象时 GPU 内存泄漏

GPU memory leakage when creating objects from sentence-transformers

r

transformer

python-3.x

pytorch

reticulate

描述

预期行为

观察到的行为

尝试求解失败

详情