时间戳

首页 > 解决方案 > 在 pytorch 中保存模型输出

nlp - 在 pytorch 中保存模型输出

问题描述

dic = [] 
for step, batch in tqdm(enumerate(train_dataloader)):
inpt = batch[0].to(device)
msks = batch[1].to(device)
#Run the sentences through the model
outputs = model_obj(inpt, msks)
dic.append( {
    'hidden_states': outputs[2],
    'pooled_output': outputs[1]})

我想在每次迭代中保存模型输出，但是对于一小组数据集，我得到了以下错误。 运行时错误：CUDA 内存不足。 请注意，如果没有以下代码，我的模型可以正常工作。

dic.append( { 'hidden_states': outputs[2], 'pooled_output': outputs[1]})

如何在每次迭代中保存这些输出？

标签： nlppytorchhuggingface-transformers

解决方案

首先，您应该始终发布完整的错误堆栈跟踪。其次，当你想存储它们以释放内存时，你应该从你的 GPU 移动输出：

dic.append( {
    'hidden_states': outputs[2].detach().cpu().tolist(),
    'pooled_output': outputs[1].detach().cpu().tolist()
})

推荐阅读