Hugging Face的BERT模型进行文本嵌入内存爆炸的解决方法

hugging face中很多预训练好的transformer模型，可以直接下载使用，节省大量时间与算力。昨天使用BERT模型进行文本嵌入。其实很简单，核心代码就几行（text是文本，batch_size是500，总共三万条文本，只取每条文本的[CLS]作文本的整体表示）：

    encoded_input = tokenizer(text[start * 500: min(start * 500 + 500, len(text))], padding=True, truncation=True, return_tensors="pt")
    output = model(**encoded_input)
    batch_embed = output[0]
    batch_embed = batch_embed.index_select(1, torch.tensor([0])).squeeze(dim=1)

然后就出现了问题，2分钟后程序就停了。我以为跑完了，将整个文本的嵌入打出来看看情况，结果显示没有此变量，再一看，jupyter内核挂了。连续跑了两次都是这样。

于是乎我用top命令查看情况：

直接被吓死。内存占用111.4，且不断在飙升，最后升到200多G，然后内核就挂了。

最后定位原因，发现是上一个batch的所有输出都被系统保留，没有被回收。然后随着batch的增多数据越来越多，内存就爆了。

使用del命令尝试删掉不需要的output，不行；使用gc.collect()强制进行垃圾回收，还是不行。最后查看博客，发现gc不一定能够回收掉，决定回不回收内存的是引用计数。于是重复进行100次垃圾回收，还是没有用（此处没有细看垃圾回收机制，迫切想解决问题，其实回收多少次都不会影响到引用计数。此外，使用del output后，打印ouput会报没有这个变量的错误，说明del成功，而垃圾回收机制没有启动。垃圾回收机制详细请看：Python垃圾回收机制详解 - Xjng - 博客园 (cnblogs.com)https://www.cnblogs.com/Xjng/p/5128269.html）

最后实在没有办法，将上述的核心代码封装成一个函数。随着函数生命周期的结束，临时变量内存被回收，至此问题解决。

def fun(start):
    encoded_input = tokenizer(text[start * 500: min(start * 500 + 500, len(text))], padding=True, truncation=True, return_tensors="pt")
    output = model(**encoded_input)
    batch_embed = output[0].index_select(1, torch.tensor([0])).squeeze(dim=1)
    joblib.dump(batch_embed, f'text_embed.pkl_{start}')

破案了，少了一句with torch.no_grad()，保存了大量梯度数据。。。。。。

添加之后，不使用函数进行封装，内存也没有爆炸。

果然菜是原罪。。。。

——————————更新——————————

原本是想要用嵌入好的文本进行下游的标签嵌入的。分批次嵌入好文本，存起来，下游任务直接从硬盘中读相当于固定住了BERT的参数，从而节省了训练模型所需的内存。但是，固定参数的方法结果奇差，评价分数是不改之前的一半。。。

但是不固定参数，几万条文本一次性进行嵌入，内存铁定吃不消，看来得换个方法了。。。