如何解决python训练数据内存泄露问题？

2021-07-05 03:07

阅读：648

标签：

Python语言越来越广泛的应用于机器学习/ 深度学习领域，是目前最火的该领域编程语言，各大深度学习框架基本都支持 python 接口。

在TensorFlow 训练模型的过程中，一般数据加载分情景各有不同。

1. 当数据量能直接全部载入内存时，当然最是方便，直接全部载入内存，然后训练即可。

2. 数据无法全部载入内存时，有多种方法。介绍其中 2 种用的多的，其一是边训练边读取（可以用多线程进行优化），其二是随机打乱所有训练数据的索引然后随机选择部分数据进行训练测试。

第2 类情景中，有可能由于对 Python 的 GC 机制理解不深（ Me… ），出现内存泄露。比如下面的第 12 行的列表就无法释放，导致内存泄露，而 11 行则不会出现内存泄露。

import numpy as np import gc import psutil import os

def get_memory_usage():

process = psutil.Process(os.getpid())

return process.memory_info().rss/(1024*1024)

print ’before load, memory: %d MB’ % get_memory_usage()#dataset = [np.zeros((6000,6000),dtype=np.float32) for _ in range(30000)]

dataset = [(np.zeros((6000,6000),dtype=np.float32),np.zeros((200,200),dtype=np.int32)) for _ in range(10000)] print ’after load, memory: %d MB’ % get_memory_usage()

print ’before release, memory: %d MB’ % get_memory_usage() del dataset

dataset = None

gc.collect() print ’after release, memory: %d MB’ % get_memory_usage()

第11 行内存释放结果：

before load, memory: 25 MB after load, memory: 220 MB before release, memory: 220 MB after release, memory: 27 MB

第12 行内存释放结果：

before load, memory: 23 MB after load, memory: 1406 MB before release, memory: 1406 MB after release, memory: 1170 MB

可以看到第12 行内存泄露严重。出现这种内存泄露时，模型训练基本只能训练几个 Epoch 就会出现 OOM 。

大力出奇迹

这种问题当然最优的是从根源解决，即研究清楚Python 的 GC 机制，解决内存泄露。

在时间比较紧或者问题比较麻烦时，有时粗暴快速的方法也不失为一种选择。