更好的python代码来减少内存的使用？

Question

我有一个大约1,900万行的数据帧，其中4个变量是纬度和经度。我创建了一个函数，用python 哈沃辛包计算纬度和经度的距离。

# function to calculate distance of 2 coordinates

def measure_distance(lat_1, long_1, lat_2, long_2):

    coordinate_start = list(zip(lat_1, long_1))
    coodrinate_end = list(zip(lat_2, long_2))
    
    distance = haversine_vector(coordinate_start, coodrinate_end, Unit.KILOMETERS)

    return distance

我使用魔术命令%%memit来测量内存的使用情况来执行计算。平均而言，内存使用量在8-10 GB之间。我在Google上运行我的工作，它有12 my的RAM，因此，有时操作会达到运行时的极限并重新启动。

%%memit

measure_distance(df.station_latitude_start.values, 
                 df.station_longitude_start.values, 
                 df.station_latitude_end.values, 
                 df.station_longitude_end.values)

peak memory: 7981.16 MiB, increment: 5312.66 MiB

有办法优化我的代码吗？

Answer 1

TL;DR:使用Numpy并按块计算结果。

对于较大的输入大小，CPython解释器需要占用的内存量。

实际上，CPython使用引用将值存储在列表中。在64位系统上，引用需要8个字节，而基本类型(浮点数和小整数)通常需要32个字节。两个浮点数的元组是一个复杂的类型，它包含元组的大小以及两个浮点数的引用(而不是值本身)。它的大小应该接近64个字节。由于您有两个列表，包含1900万对浮点数(引用)，4个列表包含1900万个浮点数(引用)，由此产生的内存应该是关于4*19e6*(8+32) + 2*19e6*(8+64) = 5.7 GB的。更别提Haversine可以复制一些内部副本了，结果也占用了一些空间。

如果您想减少内存使用量，那么使用Numpy。实际上，float数组以更紧凑的方式存储值(没有引用，没有内部标记)。可以用N x 2 Numpy 2D数组替换元组列表。结果的大小应该是关于4*19e6*8 + 2*19e6*(8*2) = 1.2 GB。此外，计算速度要快得多，Haversine在内部使用Numpy。下面是一个示例：

import numpy as np
 
# Assume lat_1, long_1, lat_2 and long_2 are of type np.array.
# Use np.array(yourList) if you want to convert it.
def measure_distance(lat_1, long_1, lat_2, long_2):
    coordinate_start = np.column_stack((lat_1, long_1))
    coordinate_end = np.column_stack((lat_2, long_2))
    return haversine_vector(coordinate_start, coordinate_end, Unit.KILOMETERS)

上面的代码大约比快25倍，。

如果您想要减少更多的内存使用量，可以通过块(例如32K值)计算坐标，然后将输出块连接起来。如果您不太关心计算距离的精度，也可以使用单精度数字而不是双精度。

下面是一个按块计算结果的示例：

def better_measure_distance(lat_1, long_1, lat_2, long_2):
    chunckSize = 65536
    result = np.zeros(len(lat_1))
    for i in range(0, len(lat_1), chunckSize):
        coordinate_start = np.column_stack((lat_1[i:i+chunckSize], long_1[i:i+chunckSize]))
        coordinate_end = np.column_stack((lat_2[i:i+chunckSize], long_2[i:i+chunckSize]))
        result[i:i+chunckSize] = haversine_vector(coordinate_start, coordinate_end, Unit.KILOMETERS)
    return result

在我的机器上，使用双精度，上面的代码大约需要800 MB，而初始实现则需要8 GB。因此，内存比少10倍！它仍然是快23倍的！使用简单的精度，上面的代码大约需要500 MB，所以内存少了16倍，而且它比快48倍！