释放PyTorch模型占用的内存并训练多个模型


我在使用CRF和Bi-LSTM-CRF训练不同数据的时候,需要在每次训练结束后,使用新的数据和模型。如果不释放之前模型占用的内容,会报错。这里的建议只在使用CPU的情况下,没有在GPU环境下测试过。

TL;DR

使用del modelgc.collect()

代码


import gc 

def crf_pipeline(data_paths, glod_data_path):
    # read glod data
    sents, glod_labels = load_data_and_labels(glod_data_path)

    for data_path in data_paths:
        # read tagged data
        tag_sents, tag_labels = load_data_and_labels(data_path)
        train_word_lists, train_tag_lists, test_word_lists, test_tag_lists = split_data(tag_sents, glod_labels, tag_labels, dev=False)

        data_path = Path(data_path)
        print("正在训练评估CRF模型:", data_path.stem)
        print('trian data: {}, test data: {}'.format(len(train_tag_lists), len(test_tag_lists)))
        crf_pred = crf_train_eval(
            (train_word_lists, train_tag_lists),
            (test_word_lists, test_tag_lists)
        )

        del crf_pred # 删除模型
        gc.collect() # 回收内存

欢迎订阅我的博客:RSS feed
知乎: 赤乐君
Blog: BrambleXu
GitHub: BrambleXu
Medium: BrambleXu

参考资料


文章作者: BrambleXu
版权声明: 本博客所有文章除特別声明外,均采用 CC BY-NC-SA 4.0 许可协议。转载请注明来源 BrambleXu !
评论
  目录