PyTorch中MLP前向传播耗时为何与总数据集大小相关？

核心原因 #

• 你观察到的耗时增长根本不是MLP前向传播导致的，而是random.sample本身的性能特性：当传入的总索引列表长度n增大时，random.sample的执行时间会线性上升。
• 因为random.sample处理大列表时，内部需要做更多重复检查（避免采样重复元素），底层随机选择逻辑会随样本池规模变大增加计算开销——哪怕最终只采batchSize个元素，样本池越大，这个步骤的耗时就越高。

验证思路 #

• 单独对random.sample步骤做耗时测试：固定batchSize，只改变n的大小，记录每次采样的耗时，就能看到n和采样耗时的正相关关系。
• 你删除采样步骤后耗时与n无关的测试结果，已经直接佐证：前向传播本身确实只和batchSize有关，MLP是被冤枉的。

优化方案 #

• 改用numpy.random.choice或者PyTorch的torch.randperm实现无重复采样，这两个方法处理大规模样本池时性能远优于random.sample：
  • PyTorch实现示例：

    # 生成随机排列的索引，取前batchSize个
    indices = torch.randperm(n)[:batchSize]
    batch_data = input_data[indices]
    

  • numpy实现示例：

    import numpy as np
    indices = np.random.choice(n, size=batchSize, replace=False)
    batch_data = input_data[indices]
    

• 若场景允许非严格无重复采样，可直接用随机整数生成，性能更优：

  indices = torch.randint(0, n, (batchSize,))
  

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