为什么PyTorch中堆叠小张量列表比更新大张量切片速度更快？

• 拷贝操作时机差异：stack_memory的循环里append操作仅存储张量的引用，不会产生任何内存拷贝开销，所有数据拷贝只在最后调用torch.stack时集中执行1次。而reserve_memory的循环中每次迭代都要执行一次out_state[:,step] = state的内存写入操作，50次迭代就对应50次拷贝，累计的拷贝开销本身就更高。
• CUDA kernel启动开销：每次切片赋值操作都会单独启动一次CUDA kernel，50次循环就会产生50次kernel启动的固定开销。而torch.stack是PyTorch原生优化的批量算子，仅需要启动1次合并处理的kernel，大幅降低了kernel启动的累计开销。
• 内存访问效率差异：你提前分配的out_state维度顺序是(batch_size, time_steps, channels, width, height)，time_steps位于第二维，每次对[:,step]赋值时是对跨步长的非连续内存块写入，GPU访存对齐率低，带宽利用率差。而torch.stack会直接将每个连续的小张量写入目标内存的连续块，访存效率更高。
• 算子优化程度差异：torch.stack作为核心内置算子，底层做了大量硬件适配优化，包括异步拷贝、批量操作合并等。而Python层面循环的逐切片赋值无法被PyTorch的算子融合优化捕获，只能每次单独处理赋值操作，没有批量优化的收益。

你可以尝试将预分配张量的维度顺序调整为(time_steps, batch_size, channels, width, height)，赋值时改为out_state[step] = state，这种连续内存写入的速度会比你当前的写法快不少，但仍然大概率略慢于stack实现，因为还是无法避免循环中多次kernel启动的开销。

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