pytorch张量高级索引介绍

PyTorch 中,张量索引操作可以使用高级索引(advanced indexing),其中索引可以是另一个张量。使用这种索引方式时,返回值的维度由索引张量的形状和原始张量的形状共同决定。以下是具体的规则和解释:

1. 基本概念

假设我们有一个张量 x 和索引张量 indices,我们通过 x[indices] 进行高级索引操作。

规则:

  • 索引张量的形状将决定返回值的形状。
  • 返回值的维度由索引张量的维度代替索引位置后的张量维度。

2. 示例讲解

示例 1:一维索引

复制代码
x = torch.tensor([[10, 20, 30], [40, 50, 60]])
indices = torch.tensor([0, 1])
result = x[indices]
  • x 的形状是 (2, 3)
  • indices 是一维张量,形状是 (2,)
  • 索引 x[indices] 的结果:
    • 取出 x 的第 0 行和第 1 行。
    • 返回值的形状是 (2, 3)

示例 2:多维索引

复制代码
x = torch.tensor([[10, 20, 30], [40, 50, 60]])
indices = torch.tensor([[0, 1], [1, 0]])
result = x[indices]

print(f"x.shape:{x.shape}")
print(f"index.shape:{index.shape}")
print(f"result.shape:{result.shape}")
print(result)

输出:

复制代码
x.shape:torch.Size([2, 3])
index.shape:torch.Size([2, 2])
result.shape:torch.Size([2, 2, 3])
tensor([[[10, 20, 30],
         [40, 50, 60]],

        [[40, 50, 60],
         [10, 20, 30]]])

示例 3:多维组合索引

复制代码
x = torch.tensor([[10, 20, 30], [40, 50, 60]])
rows = torch.tensor([0, 1])
cols = torch.tensor([1, 2])
result = x[rows, cols]
  • x 的形状是 (2, 3)
  • rowscols 都是一维张量,形状为 (2,)
  • 索引 x[rows, cols]
    • 分别取出 x[0, 1]x[1, 2]
    • 返回值是 (20, 60),形状为 (2,)

示例 4:广播索引

复制代码
x = torch.tensor([[10, 20, 30], [40, 50, 60]])
rows = torch.tensor([[0], [1]])
cols = torch.tensor([0, 2])
result = x[rows, cols]
  • x 的形状是 (2, 3)
  • rows 的形状是 (2, 1)cols 的形状是 (2,)
  • 索引 x[rows, cols]
    • rowscols 会广播成 (2, 2)
    • 返回值的形状是 (2, 2)

示例 5:更复杂的张量索引操作

AF3 AtomAttentionEncoder类的init_pair_repr方法解读-CSDN博客中的 张量的高级索引

总结:

  • 索引张量的形状直接决定了返回张量的形状。
  • 当多个索引张量时,它们会广播以匹配维度,然后返回广播后形状的张量。
相关推荐
凯子坚持 c19 分钟前
从 0 到 1:ComfyUI AI 工作流抠图构建全实践
大数据·人工智能
背太阳的牧羊人1 小时前
OpenAI Embedding 和密集检索(如 BERT/DPR)进行语义相似度搜索有什么区别和联系
人工智能·bert·embedding
光与电子KOYUELEC加油奋斗2 小时前
Molex莫仕连接器:增强高级驾驶辅助系统,打造更安全的汽车
人工智能·光与电子
__lost4 小时前
MATLAB画出3d的常见复杂有机分子和矿物的分子结构
开发语言·人工智能·matlab·化学·分子结构
每天都要写算法(努力版)4 小时前
【神经网络与深度学习】五折交叉验证(5-Fold Cross-Validation)
人工智能·深度学习·神经网络
hi星尘5 小时前
深度解析:基于Python的微信小程序自动化操作实现
python·微信小程序·自动化
郭不耐5 小时前
DeepSeek智能时空数据分析(六):大模型NL2SQL绘制城市之间连线
人工智能·数据分析·时序数据库·数据可视化·deepseek
Doker 多克5 小时前
Django 缓存框架
python·缓存·django
winfredzhang6 小时前
Deepseek 生成新玩法:从文本到可下载 Word 文档?思路与实践
人工智能·word·deepseek
KY_chenzhao6 小时前
ChatGPT与DeepSeek在科研论文撰写中的整体科研流程与案例解析
人工智能·机器学习·chatgpt·论文·科研·deepseek