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

Question

描述

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

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

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 缓存)

尝试解决该问题但不成功

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

详细信息

Windows 10 家庭版

Python 3.7.4

R 4.0.1

网状1.16

torch 1.3.1

tensorflow 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()")

  }

而且效果非常好。