PyTorch内存分配疑惑及无内存增长的字典Flatten方法咨询

代码实现 #

from typing import OrderedDict
import torch

def flatten_dict(
        aligned_feat: OrderedDict[str, list[torch.Tensor]], 
    ):
    num_batches = len(list(aligned_feat.values())[0])
    flat_features = [None] * num_batches

    for k in list(aligned_feat.keys()):
        for b in range(num_batches):
            if flat_features[b] is None:
                flat_features[b] = aligned_feat[k][b]
            else:
                flat_features[b] = torch.cat((flat_features[b], aligned_feat[k][b]))
        # free up mem
        del aligned_feat[k]
        print(f"{k}: {torch.cuda.max_memory_allocated()}")

    return flat_features

if __name__ == "__main__":
    keys = ['0', '1', '2', '3']
    d = OrderedDict()

    for k in list(keys):
        d[k] = [torch.rand((50, 256, 7, 7)).to("cuda") for _ in range(4)]

    print(str(torch.cuda.max_memory_allocated()))
    data = flatten_dict(d)
    print(str(torch.cuda.max_memory_allocated()))

运行输出 #

41943040 
0: 41943040 
1: 54487040 
2: 56995840 
3: 59504640 
59504640

疑问点 #

• 为何处理key '0'时内存未增长？
• 首次增长后为何持续增加？
• 是否存在另一种“展平”字典的方式，不会增加最大分配内存（函数内部也不增加）？注：输入中list[torch.Tensor]的shape[0]可能不同，但shape[1:]固定为[256,7,7]，展平后不需要原字典格式。

内存变化原因分析 #

处理key '0'时内存未增长的原因 #

处理第一个key '0'时，代码只是将aligned_feat['0'][b]直接赋值给flat_features[b]，这是引用传递，没有创建新的GPU张量，flat_features中的张量与原字典中的张量共享同一块内存，因此最大内存分配值与初始状态一致。

首次增长后持续增加的原因 #

从处理key '1'开始，每次执行torch.cat都会创建新的张量：

• 第一次拼接（key '1'）：需要同时保留原flat_features[b]（来自key '0'）和aligned_feat['1'][b]，再分配新内存存储拼接结果，内存峰值因此上升；
• 后续key '2'、'3'的拼接逻辑同理：每次拼接都要保留当前flat_features[b]和新张量，生成更大的新张量，内存峰值持续累加；
• 虽然代码删除了原字典的key，但PyTorch的CUDA内存回收并非即时（需等待垃圾回收或显式释放），且拼接过程中临时占用的内存已被计入最大分配值，因此峰值只会上升不会下降。

无额外内存增长的展平方案 #

核心思路是预分配目标内存+原地拷贝，避免torch.cat创建新张量带来的额外内存占用：

from typing import OrderedDict
import torch

def flatten_dict_no_extra_mem(aligned_feat: OrderedDict[str, list[torch.Tensor]]):
    num_batches = len(next(iter(aligned_feat.values())))
    # 计算每个batch需要的总shape[0]长度
    batch_total_dims = []
    for b in range(num_batches):
        total_dim = sum(t.shape[0] for t in [aligned_feat[k][b] for k in aligned_feat.keys()])
        batch_total_dims.append(total_dim)
    
    # 预分配每个batch的最终张量
    flat_features = []
    for b in range(num_batches):
        flat_feat = torch.empty((batch_total_dims[b], 256, 7, 7), device='cuda')
        flat_features.append(flat_feat)
    
    # 逐个key拷贝数据到预分配张量中
    current_pos = [0] * num_batches
    for k in list(aligned_feat.keys()):
        for b in range(num_batches):
            t = aligned_feat[k][b]
            dim_len = t.shape[0]
            # 原地拷贝，无新内存分配
            flat_features[b][current_pos[b]:current_pos[b]+dim_len] = t
            current_pos[b] += dim_len
        # 释放原张量内存并显式回收CUDA内存
        del aligned_feat[k]
        torch.cuda.empty_cache()
        print(f"{k}: {torch.cuda.max_memory_allocated()}")
    
    return flat_features

if __name__ == "__main__":
    keys = ['0', '1', '2', '3']
    d = OrderedDict()

    for k in keys:
        d[k] = [torch.rand((50, 256, 7, 7)).to("cuda") for _ in range(4)]

    print(str(torch.cuda.max_memory_allocated()))
    data = flatten_dict_no_extra_mem(d)
    print(str(torch.cuda.max_memory_allocated()))

方案优势 #

• 预分配内存：一次性分配最终所需的全部内存，避免拼接时的临时内存占用；
• 原地拷贝：直接将原张量数据拷贝到预分配张量，无新张量创建；
• 即时回收：删除原字典key后显式调用torch.cuda.empty_cache()，及时释放原张量占用的内存，最大内存分配值可保持与初始值接近。

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