PyTorch jagged布局嵌套张量致DataLoader性能骤降的原因与优化

问题背景 #

PyTorch官方推荐使用torch.jagged布局创建嵌套张量，但实际在Dataset中使用该布局并通过DataLoader加载时，性能远低于torch.strided布局。测试显示jagged布局的DataLoader耗时约为strided布局的14倍，以下是复现测试代码：

import torch


# dataset with nested tensor
class TestDatasetNested(torch.utils.data.Dataset):
    def __init__(self, N=100, dl=100, layout=torch.jagged):
        self.N = N
        self.dl = dl

        # set std deviations
        nnoise = torch.randint(1, high=5, size=(self.N,))
        sigmas = [torch.rand(n) for n in nnoise]

        self.sigmas = torch.nested.nested_tensor(sigmas, layout=layout)

    def __getitem__(self, i):
        sigmas = self.sigmas[i % self.N]

        return torch.cat([sigma * torch.randn((self.dl, 1)) for sigma in sigmas], dim=-1).sum(dim=-1)

    def __len__(self):
        return self.N


# create dataset with jagged and strided nested tensor layouts
dataset_jagged = TestDatasetNested(layout=torch.jagged)
dataset_strided = TestDatasetNested(layout=torch.strided)

# create dataloader for both cases
dl_jagged = torch.utils.data.DataLoader(dataset=dataset_jagged, batch_size=10)
dl_strided = torch.utils.data.DataLoader(dataset=dataset_strided, batch_size=10)

def nepochs(dl, n=100):
    for i in range(n):
        for _ in dl:
            pass

%timeit nepochs(dl_jagged)
# 4.57 s ± 17.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit nepochs(dl_strided)
# 322 ms ± 3.41 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

性能差异原因 #

• DataLoader批处理适配不足：torch.strided布局的嵌套张量本质是普通张量的包装，DataLoader对其批处理的内存对齐、批量操作逻辑已高度优化；而torch.jagged是针对不规则张量的新实现，当前DataLoader对其批处理的支持不完善，采样、拼接环节会产生额外内存拷贝与计算开销。
• 单元素访问开销高：__getitem__中访问jagged布局张量的单个元素时，需解析jagged结构的维度信息并定位内存，相比strided布局的直接内存访问，逻辑更复杂、耗时更长。
• 操作叠加放大差距：测试代码中每次__getitem__包含随机张量生成与循环拼接操作，jagged布局的访问开销会与这些操作叠加，进一步扩大性能差距。

可行优化方案 #

• 预生成固定形状张量：若业务允许，将嵌套张量填充为固定形状的普通张量（用0或特定值补全），直接利用DataLoader对普通张量的高效批处理能力，完全规避嵌套张量的性能问题。
• 自定义collate_fn优化批处理：针对jagged布局实现自定义collate_fn，手动优化批量拼接逻辑，减少不必要的内存操作。示例代码如下：

def jagged_collate_fn(batch):
    # 自定义批处理逻辑，直接将batch结果整理为jagged张量
    return torch.nested.nested_tensor(batch, layout=torch.jagged)

# 使用优化后的collate_fn创建DataLoader
dl_jagged_opt = torch.utils.data.DataLoader(
    dataset=dataset_jagged, 
    batch_size=10, 
    collate_fn=jagged_collate_fn
)

• 延迟嵌套张量创建：不在Dataset的__init__中预先构建jagged张量，而是在__getitem__中直接处理原始列表数据，仅在最终批量处理时再转换为jagged布局，避免提前创建带来的索引开销。
• 升级PyTorch版本：PyTorch对jagged布局的支持持续迭代，新版本可能修复了DataLoader交互中的性能瓶颈，尝试升级至最新稳定版测试性能。

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