PyTorch模型训练过程内存泄漏问题解决

近日,在模型训练过程中,发现过一段时间后进程会被kill,观察发现是由于内存泄漏问题造成的。通过逐行代码注释,发现问题在于数据集中的此行代码:

info = self.data_list[index]

这里,self.data_list是dataset类中保存数据字典的列表,index是列表索引。在模型训练的全部过程中,都使用该字典进行数据在不同模型和模块中的传递。由于数据字典始终被dataset所引用,故在全程序运行过程中会始终存在,包括不断加入字典的新数据,因此需要此行代码为

info = copy.deepcopy(self.data_list[index]).

你可能感兴趣的:(pytorch,人工智能,python,内存泄漏)