Python：如何高效提取多维张量中第二维度的每隔n个元素

对于形状为(12, 19601, 1000)的张量，提取第二维度每隔n个元素的最优方式是利用张量框架原生的切片操作——这是底层硬件加速实现的，完全避免了Python层面的循环，是处理大规模张量时速度最快的方案。

PyTorch 实现 #

直接使用切片语法，针对第二维度指定步长为n：

import torch

# 假设你的张量为tensor
tensor = torch.randn(12, 19601, 1000)
n = 5  # 示例步长

# 提取第二维度每隔n个元素
result = tensor[:, ::n, :]

如果需要严格保证结果第二维度长度为19601 / n（即要求19601能被n整除），可先截断第二维度到能被n整除的长度，再切片：

# 计算截断后的第二维度长度
trunc_len = (tensor.shape[1] // n) * n
# 截断后再切片
result = tensor[:, :trunc_len:n, :]

TensorFlow 实现 #

同样支持原生切片语法，或使用tf.gather指定索引范围（效果一致）：

import tensorflow as tf

# 假设你的张量为tensor
tensor = tf.random.normal((12, 19601, 1000))
n = 5  # 示例步长

# 方法1：切片语法
result = tensor[:, ::n, :]

# 方法2：tf.gather指定步长索引
indices = tf.range(0, tensor.shape[1], n)
result = tf.gather(tensor, indices, axis=1)

效率说明 #

上述方法的核心优势在于：

• 切片操作是框架底层优化的，会直接调用CPU向量指令或CUDA核函数执行，无额外内存拷贝开销。
• 无论是CPU还是GPU环境，都能充分利用硬件并行能力，处理大规模张量时性能远优于手动循环或其他间接方法。

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