内存优化——使用pandas读取千万级数据

超级干货：Python优化之使用pandas读取千万级数据

环境：Linux-cenos5

processor : 31

model : 62

model name : Intel(R) Xeon(R) CPU E5-2640 v2 @ 2.00GHz

cpu MHz : 2000.066

cache size : 20480 KB

memory : 125G

在如上所述的单机环境中，使用一些优化可以使基于pandas数据格式的模型训练数据容量由600W增长为至少2000W，训练时间减少为1/5。具体方案如下：

数据读取优化

数据量4200W行，193列，每列存储为string类型的单精度浮点数，文件表由csv格式存储，总大小16GB+。通过如下语句读取到dataframe中去：

1 2	`df_train` `=` `pd.read_csv(path,header=None,sep=',',nrows=40000000,error_bad_lines=False,delimiter="\t",lineterminator="\n",keep_default_na=True)`

　　经过测试，当nrows读取行数超过800W条时，df_train占内存超过80G，在后续的步骤中涉及到切片和数据集复制时会直接崩溃，超过1200W条时会直接无法读取。首先考虑优化读取方式：

na_vals = ["\\N"," ","","NULL"]

df_tmp = []

df_train = pd.DataFrame(index=["0"],dtype=np.float32)

count = 0

for chunk in pd.read_csv(path,header=None,sep=',',chunksize=200000,nrows=10000000,error_bad_lines=False,delimiter="\t",lineterminator="\n",keep_default_na=True, na_values=na_vals):

df_tmp.append(chunk[1:])

del chunk

print("the chunk " + str(count) + " has been stored...")

print("the mem-cost is now: ", str(sys.getsizeof(df_tmp)/(1) ), "MB \n")

可以分为以下几个点解释

使用na_values进行读取时的空值替换，代替了读取全部数据后的空值replace方式替换，减少了后续的处理时间
使用chunksize进行数据读取，read_csv()方法中有一个参数chunksize可以指定一个CHUNKSIZE分块大小来读取文件。与直接使用DF进行遍历不同的是，此时它是一个TextFileReader类型的对象。

通过循环每次读取分块数据，再通过list拼接起来。

为什么不直接用concat进行迭代拼接？

因为concat只能使用值语义赋值，每次concat时都会先创建一个临时副本，再赋值给对应变量，这个过程中当数据量很大时，内存占用在赋值完成，删除原有dataframe前会造成极高的内存占用峰值。而python的list可以进行连续内存块的直接追加，不会产生额外的内存开销。

2. 每次读取完一个chunk，都进行删除，释放对应内存。

数据转换优化

经过读取优化后，脚本最多可读取超过2000W的数据量，是之前读取数据量的三倍以上。但是此时读取完成后，在后续的处理中仍然会产生内存报错，首先报错的就是

df_X = pd.DataFrame(data=df_tmp, dtype=np.float)

内存报错原因：

当我们将读取的数据list经由pandas的构造函数进行类型转换时，会产生两个超大的dataframe/list，一个是等号右边的临时表，一个是df_tmp本身，此时内存超过限制，程序崩溃。

经过分析采用如下方案可避免此部分报错

idx = 0

for i in range(len(df_tmp)):

tmp = pd.DataFrame(data=df_tmp[idx], dtype=np.float)

df_train = pd.concat([df_train,tmp], ignore_index=True)

del df_tmp[idx],tmp

print(i)

print("the remaining chunk is: ", count)

print("the frame size is: ", df_train.memory_usage().sum() / (1024 ** 2), "MB") count -= 1

同样通过循环，每次只处理df_tmp这个存储了所有数据的临时list中的一个chunk，并在数据转换期间进行迭代拼接，每次增量更新df_train主表。在更新完成后立即释放该list中对应的内存。经过优化，此部分内存占用减半

数据类型优化

pandas中的许多数据类型具有多个子类型，它们可以使用较少的字节去表示不同数据，比如，float型就有float16、float32和float64这些子类型。这些类型名称的数字部分表明了这种类型使用了多少比特来表示数据，比如刚才列出的子类型分别使用了2、4、8个字节。下面这张表列出了pandas中常用类型的子类型：

一个int8类型的数据使用1个字节（8位比特）存储一个值，可以表示256（2^8）个二进制数值。这意味着我们可以用这种子类型去表示从-128到127（包括0）的数值。原始数据表基本都是由浮点数组成，因此可以考虑使用数据类型进行优化。

tmp0 = pd.DataFrame(data=df_tmp[idx])

tmp1 = pd.DataFrame(data=df_tmp[idx], dtype=object)

tmp2 = pd.DataFrame(data=df_tmp[idx], dtype=np.float)

tmp3 = pd.DataFrame(data=df_tmp[idx], dtype=np.float32)

# print("the remaining chunk is: ", count)

print("the frame size is: ", df_train.memory_usage().sum() / (1024 ** 2), "MB")

数据量为2W的情况下，pandas内存占用情况如下(默认类型为str)

一开始的代码中没有进行类型的转换，因此默认根据csv的格式进行初始化，所以Dataframe的数据类型全都是str，在读取之后的存储开销非常大，16G的csv读取进来后直接变成120G。使用数据类型优化后，结合之前的分块迭代读取，2000W的数据(8G)读取由60G的内存占用减少为不到15G，且省去了后续直接的类型转换。

小技巧

在每次使用完某个对象后，如果其在后续的建模过程中没有别的用处，可以手动删除，或者删除变量引用，触发GC机在读入数据集的时候指定列的最优数据类型。pandas.read_csv()函数有一些参数可以做到这一点。dtype参数接受一个以列名（string型）为键字典、以Numpy类型对象为值的字典

dict = {x:np.float16 for x in col}

df_train = pd.read_csv(path,header=None,sep=',',error_bad_lines=False,delimiter="\t",lineterminator="\n",

keep_default_na=True, converters=dict, na_vals=na_vals)

　　可以考虑使用开源的pandas分布式机器学习库dask，在单机环境下运行即可