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

Question

我有一个大约 1900 万行的数据框，其中 4 个变量是纬度和经度。我在 python haversine 包的帮助下创建了一个计算纬度和经度距离的函数。

# 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 之间。我在具有 12GB RAM 的 Google Colab 上运行我的工作，结果，有时操作达到运行时限制并重新启动。

%%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 并通过 chunk 计算结果。

CPython 解释器占用的内存量与大输入大小有关。

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

如果你想减少内存使用，那么使用 Numpy。实际上，浮点 Numpy 数组以更紧凑的方式存储值（没有引用，没有内部标记）。您可以用N x 2 Numpy 二维数组替换元组列表。结果大小应约为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

在我的机器上，使用 double 精度，上面的代码大约需要 800 MB，而初始实现需要 8 GB。因此，内存减少了 10 倍！它仍然快了 23 倍！使用 simple 精度，上面的代码大约需要 500 MB，所以 内存减少了 16 倍，速度提高了 48 倍！