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()，保存了大量梯度数据。。。。。。
添加之后，不使用函数进行封装，内存也没有爆炸。
果然菜是原罪。。。。