如何降低Python程序运行时的内存占用（含GPU内存优化）

问题描述 #

运行文件读取类Python程序时遭遇内存占用异常：待读取文件仅100MB，但进程内存占用高达1.6GB（仅导入必要库，未定义额外变量）。将变量迁移至GPU后，GPU内存占用也达到600MB，求可行的内存优化方案。

相关代码 #

import numpy as np
import torch
import time
import struct

if __name__ == "__main__":
    graphEdge = []
    boundList = []

    file_path = "./../../srcList.bin"
    with open(file_path, 'rb') as file:
        while True:
            data = file.read(4)
            if not data:
                break
            integer = struct.unpack('i', data)[0]
            graphEdge.append(integer)


    file_path = "./../../range.bin"
    with open(file_path, 'rb') as file:
        while True:
            data = file.read(4)
            if not data:
                break
            integer = struct.unpack('i', data)[0]
            boundList.append(integer)

    graphEdge = torch.Tensor(graphEdge).to(torch.int).to('cuda:0')
    boundList = torch.Tensor(boundList).to(torch.int).to('cuda:0')
    memory_size = graphEdge.element_size() * graphEdge.numel()
    print(f"Tensor memory: {memory_size/(1024*1024)} MB")

优化方案 #

1. 跳过Python列表，直接用numpy/torch读取二进制文件 #

Python列表存储单个int对象的开销远大于原始4字节整数（64位系统下每个int约占28字节），这是内存暴涨的核心原因。直接用numpy或torch的二进制读取接口，一次性将文件加载为连续内存的数组/张量，避免中间列表的额外开销：

import numpy as np
import torch

# 直接读取二进制文件为numpy int32数组
graphEdge_np = np.fromfile("./../../srcList.bin", dtype=np.int32)
boundList_np = np.fromfile("./../../range.bin", dtype=np.int32)

# 转换为GPU张量，指定正确数据类型
graphEdge = torch.from_numpy(graphEdge_np).to('cuda:0', dtype=torch.int32)
boundList = torch.from_numpy(boundList_np).to('cuda:0', dtype=torch.int32)

# 可选：删除numpy数组释放CPU内存
del graphEdge_np, boundList_np

2. 避免不必要的数据类型转换 #

原代码中torch.Tensor(graphEdge).to(torch.int)会先将列表转为默认的float32张量，再转换为int类型，中间会创建一个临时float张量，额外占用一倍内存。直接指定 dtype 创建张量：

# 替代原转换逻辑，直接生成int32张量
graphEdge = torch.tensor(graphEdge, dtype=torch.int32).to('cuda:0')

3. 分批次读取处理（按需） #

如果后续业务逻辑不需要一次性加载全部数据，可分批次读取处理，读完一批释放一批内存：

batch_size = 1000000  # 每批次读取100万条数据
with open("./../../srcList.bin", 'rb') as file:
    while True:
        # 读取一批4字节整数的二进制数据
        batch_data = file.read(batch_size * 4)
        if not batch_data:
            break
        # 直接从二进制缓冲区创建GPU张量
        batch_tensor = torch.frombuffer(batch_data, dtype=torch.int32).to('cuda:0')
        # 执行批次处理逻辑
        # ...
        # 手动释放当前批次内存
        del batch_tensor
        torch.cuda.empty_cache()

4. 及时释放无用对象 #

转换为张量后，手动删除原始Python列表，并清理GPU缓存，避免内存占用冗余：

graphEdge = torch.tensor(graphEdge, dtype=torch.int32).to('cuda:0')
del graphEdge  # 删除原列表，触发GC回收CPU内存
torch.cuda.empty_cache()  # 清理GPU未使用的缓存内存

5. 排查内存细节 #

• CPU内存：使用memory_profiler库逐行分析内存占用，定位异常点；
• GPU内存：执行torch.cuda.memory_summary()查看详细内存分配，确认是否有未释放的临时对象。

内容的提问来源于stack exchange，提问作者YA xiang