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

Posted 2023-03-05

【中文标题】更好的python代码来减少内存使用？

【英文标题】：Better python code for reduce memory usage?

【发布时间】：2021-03-13 13:51:35

【问题描述】：

我有一个大约 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

有没有办法优化我的代码？

【问题讨论】：

您的输入数据帧使用了多少内存（维度和字节），您预计输出距离数组会占用多少内存？您的数据集可能太大，您的机器无法在计算期间复制数据。 list(zip... 行肯定会复制您的数据，haversine_vector 函数可能会生成临时中间数组。

与.info()，我可以看出数据框正在使用大约 600 MB 内存。我没有具体的期望，只是想在看到 Google Colab 崩溃时找出该命令的内存使用情况。我尝试创建coordinate_start 和coordinate_end 然后在haversine 中使用它们，每个操作都使用8GB 的​​内存。因此，我决定将所有内容都包装在函数中，看看有什么不同。最后，内存使用情况几乎相同，有/没有功能。我的目标是减少内存使用量，这样 Google Colab 就不会崩溃。

600 MB 固然很大，但 12 GB 可以提供 20 个副本（远低于此值，因为内存还用于其他用途）。我怀疑haversine_vector 只是制作了太多临时数组。如果您不运行haversine_vector 函数，请尝试查看使用了多少内存（通过在前面键入# 将其注释掉）。文档建议它返回一个与输入长度相同的 numpy 距离数组，因此您可以返回一个虚拟值 distance = numpy.zeros(len(coordinate_start))。如果内存使用量急剧下降，那么我们可以确定是haversine_vector 的错。

【参考方案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 倍！

【讨论】：

TQ @Jérôme Richard 为您提供建议。使用np.column_stack，优化运算peak memory: 2806.70 MiB, increment: 0.02 MiB，5秒内完成计算。真是太神奇了！！能否用示例代码进一步详细说明chunk 方法？

我更新了答案以添加有关此方法和我的结果的更多信息；）。

再次感谢。我会花一些时间研究这个方法。