如何用切片操作移除PyTorch嵌套张量中每个子张量的中间10个元素？

给定如下输入张量：

import torch

input_tensor = torch.tensor([[ 0.1585, -0.1414, -0.0166,  0.1008, -0.0693,  0.1748,  0.1587, -0.0262,
         -0.1192,  0.0404,  0.0661, -0.1363, -0.1636,  0.0900,  0.1460, -0.1348,
          0.1293,  0.0736,  0.1186,  0.0339, -0.1225,  0.1066,  0.1314, -0.0322,
         -0.1445,  0.0966, -0.1792,  0.0682,  0.0310,  0.1446],
        [ 0.0143, -0.0800,  0.0851,  0.1199, -0.1364, -0.0677, -0.0623, -0.0921,
          0.0284, -0.0700, -0.1295, -0.1681, -0.1670,  0.0758, -0.1469,  0.0280,
          0.0460, -0.1545, -0.1377,  0.1204, -0.0134, -0.0046,  0.1248, -0.1202,
         -0.1177, -0.0598,  0.1648,  0.0955,  0.1262, -0.0785]])

需要移除每个子张量的中间10个元素，得到如下输出：

tensor([[ 0.1585, -0.1414, -0.0166,  0.1008, -0.0693,  0.1748,  0.1587, -0.0262,
         -0.1192,  0.0404, -0.1225,  0.1066,  0.1314, -0.0322,
         -0.1445,  0.0966, -0.1792,  0.0682,  0.0310,  0.1446],
        [ 0.0143, -0.0800,  0.0851,  0.1199, -0.1364, -0.0677, -0.0623, -0.0921,
          0.0284, -0.0700, -0.0134, -0.0046,  0.1248, -0.1202,
         -0.1177, -0.0598,  0.1648,  0.0955,  0.1262, -0.0785]])

解决方案：利用切片拼接 #

观察输入张量，每个子张量长度为30，中间10个元素对应索引10:20（Python左闭右开切片规则）。可以通过切片获取前后两段元素，再用torch.cat拼接实现需求，代码简洁高效：

# 取每个子张量的前10个元素和后10个元素，沿维度1拼接
output_tensor = torch.cat([input_tensor[:, :10], input_tensor[:, 20:]], dim=1)

print(output_tensor)

运行上述代码后，输出结果与期望完全一致。

通用化扩展 #

如果子张量长度不固定，可先计算中间10个元素的起始和结束索引：

seq_len = input_tensor.size(1)
start_idx = (seq_len - 10) // 2
end_idx = start_idx + 10

output_tensor = torch.cat([input_tensor[:, :start_idx], input_tensor[:, end_idx:]], dim=1)

这样无论子张量长度是多少，都能精准移除中间10个元素。

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