用于减少内存使用的更好的 Python 代码？

问题描述

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

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

解决方法

TL;DR：使用 Numpy 并通过 chunk 计算结果。

CPython 解释器占用的内存量在大输入大小的情况下是预期的。

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

如果您想减少内存使用量，请使用 Numpy。实际上，float Numpy 数组以更紧凑的方式存储值（没有引用，没有内部标记）。您可以用 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,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_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,Unit.KILOMETERS)
    return result

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

google-colaboratory haversine performance performance performance python