在 pytorch 中保存模型输出答案

【问题标题】：save model output in pytorch在 pytorch 中保存模型输出
【发布时间】：2021-07-15 15:37:52
【问题描述】：

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]})

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

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

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

【问题讨论】：

标签： nlp pytorch huggingface-transformers

【解决方案1】：

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

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

【讨论】：