大佬的原版请参考代码里面的链接,下面显示的是我的修改版
#减少pandas占用的内存
# @from: https://www.kaggle.com/arjanso/reducing-dataframe-memory-size-by-65/code
# @liscense: Apache 2.0
# @author: weijian
'''函数说明
仅支持object格式和整形,浮点数
入参:
props:DataFrame
floattype:浮点数转换类型(所有的浮点数都会被转换到该类型),默认np.float16
floattoint:是否强制将float转换成int,默认为True
floatthreshold:float转换成int阈值,大于等于这个阈值则不会转换,默认为0.01
fillnavalue:缺失值填充值,默认-999
'''
def reduce_mem_usage(props,floattype = np.float16, floattoint = True, floatthreshold = 0.01,fillnavalue = -999):
# 计算当前内存
start_mem_usg = props.memory_usage().sum() / 1024 ** 2
print("Memory usage of the dataframe is :", start_mem_usg, "MB")
# 哪些列包含空值,空值用-999填充。why:因为np.nan当做float处理
#NAlist = []
for col in props.columns:
# 这里只过滤了objectd格式,如果你的代码中还包含其他类型,请一并过滤
if (props[col].dtypes != object):
# 判断是否是int类型
isInt = False
mmax = props[col].max()
mmin = props[col].min()
# Integer does not support NA, therefore Na needs to be filled
if not np.isfinite(props[col]).all():
#NAlist.append(col)
props[col].fillna(fillnavalue, inplace=True) # 用-999填充
# test if column can be converted to an integer
asint = props[col].fillna(0).astype(np.int64)
result = np.fabs(props[col] - asint)
result = result.sum()
if result < floatthreshold: # 绝对误差和小于floatthreshold认为可以转换的,要根据task修改
if result == 0:
isInt = True
else:
isInt = floattoint
# make interger / unsigned Integer datatypes
if isInt:
if mmin >= 0: # 最小值大于0,转换成无符号整型
if mmax <= 255:
props[col] = props[col].astype(np.uint8)
elif mmax <= 65535:
props[col] = props[col].astype(np.uint16)
elif mmax <= 4294967295:
props[col] = props[col].astype(np.uint32)
else:
props[col] = props[col].astype(np.uint64)
else: # 转换成有符号整型
if mmin > np.iinfo(np.int8).min and mmax < np.iinfo(np.int8).max:
props[col] = props[col].astype(np.int8)
elif mmin > np.iinfo(np.int16).min and mmax < np.iinfo(np.int16).max:
props[col] = props[col].astype(np.int16)
elif mmin > np.iinfo(np.int32).min and mmax < np.iinfo(np.int32).max:
props[col] = props[col].astype(np.int32)
elif mmin > np.iinfo(np.int64).min and mmax < np.iinfo(np.int64).max:
props[col] = props[col].astype(np.int64)
else: # 注意:这里对于float都转换成floattype,需要根据你的情况自己更改
props[col] = props[col].astype(floattype)
# print("dtype after", props[col].dtype)
# print("********************************")
print("___MEMORY USAGE AFTER COMPLETION:___")
mem_usg = props.memory_usage().sum() / 1024**2
print("Memory usage is: ",mem_usg," MB")
print("This is ",100*mem_usg/start_mem_usg,"% of the initial size")
return props