【pytorch】pytorch中的高级索引

这里只介绍pytorch的高级索引,是一些奇怪的切片索引

基本版

a[[0, 2], [1, 2]] 等价 a[0, 1] 和 a[2, 2],相当于索引张量的第一行的第二列和第三行的第三列元素;

a[[1, 0, 2], [0]] 等价 a[1, 0] 和 a[0, 0] 和 a[2, 0],相当于索引张量的第二行第一列的元素、张量第一行和第一列的元素以及张量第三行和第一列的元素

python 复制代码
import torch
a = torch.arange(9).view([3, 3])

print(a)
b = a[[0, 2], [1, 2]]

print(b)

c = a[[1, 0, 2], [0]]

print(c)


# ---------output----------
# tensor([[0, 1, 2],
#         [3, 4, 5],
#         [6, 7, 8]])
# tensor([1, 8])
# tensor([3, 0, 6])

# 这里参考了:https://zhuanlan.zhihu.com/p/509591863

高级索引的原则:索引中有: 就代表着改维度全部取,在哪个维度放置索引,就代表想取哪个维度的内容

扩展A:

python 复制代码
import torch
a = torch.arange(30).view([2, 5, 3])  # 假如a代表N, x, y
print(a)

b = torch.tensor([[0, 2],
                  [1, 2]])
print(a[:, b[0, :], :])

# ------output-------
# tensor([[[ 0,  1,  2],
#          [ 3,  4,  5],
#          [ 6,  7,  8],
#          [ 9, 10, 11],
#          [12, 13, 14]],
#
#         [[15, 16, 17],
#          [18, 19, 20],
#          [21, 22, 23],
#          [24, 25, 26],
#          [27, 28, 29]]])
# tensor([[[ 0,  1,  2],
#          [ 6,  7,  8]],
#
#         [[15, 16, 17],
#          [21, 22, 23]]])

# 上述代码含义是对a的所有N, 按b中的第一行取出所有的行

扩展B:

python 复制代码
a = torch.arange(30).view([2, 5, 3])
print(a)

b = torch.tensor([[0, 2],
                  [1, 2]])
print(a[:, b[:, 0], b[:, 1]])

# ----------output--------------
# tensor([[[ 0,  1,  2],
#          [ 3,  4,  5],
#          [ 6,  7,  8],
#          [ 9, 10, 11],
#          [12, 13, 14]],
# 
#         [[15, 16, 17],
#          [18, 19, 20],
#          [21, 22, 23],
#          [24, 25, 26],
#          [27, 28, 29]]])
# tensor([[ 2,  5],
#         [17, 20]])

# 上述代码含义是对a的所有batch, 按b中的元素取出a中的x, y; 取N次

扩展C: (最抽象的一次)

python 复制代码
a = torch.arange(30).view([2, 5, 3])
print(a)

b = torch.tensor([[0, 2],
                  [1, 2]])
print(a[:, b, :])

# ------output-------
# tensor([[[ 0,  1,  2],
#          [ 3,  4,  5],
#          [ 6,  7,  8],
#          [ 9, 10, 11],
#          [12, 13, 14]],
# 
#         [[15, 16, 17],
#          [18, 19, 20],
#          [21, 22, 23],
#          [24, 25, 26],
#          [27, 28, 29]]])
# tensor([[[[ 0,  1,  2],
#           [ 6,  7,  8]],
# 
#          [[ 3,  4,  5],
#           [ 6,  7,  8]]],
# 
# 
#         [[[15, 16, 17],
#           [21, 22, 23]],
# 
#          [[18, 19, 20],
#           [21, 22, 23]]]])

# 上述代码含义是对a的所有batch, 按b中的元素取出a中行; 取N * b[0]次

torch.gather函数

本来想使用torch.gather函数完成上述功能,实验后发现并不直观,还是用高级索引吧。这里放个torch.gather函数单独的内容吧。

python 复制代码
import torch

tensor_0 = torch.arange(3, 12).view(3, 3)
print(tensor_0)

index = torch.tensor([[2, 1, 0]])
tensor_1 = tensor_0.gather(0, index)
print(tensor_1)


index = torch.tensor([[2, 1, 0]])
tensor_1 = tensor_0.gather(1, index)
print(tensor_1)

#-------------output------------
# tensor([[ 3,  4,  5],
#         [ 6,  7,  8],
#         [ 9, 10, 11]])
# tensor([[9, 7, 5]])
# tensor([[5, 4, 3]])

# torch.gather的理解
# index=[ [x1,x2,x2],
# [y1,y2,y2],
# [z1,z2,z3] ]
# 
# 如果dim=0
# 填入方式
# [ [(x1,0),(x2,1),(x3,2)]
# [(y1,0),(y2,1),(y3,2)]
# [(z1,0),(z2,1),(z3,2)] ]
#
# 如果dim=1
# [ [(0,x1),(0,x2),(0,x3)]
# [(1,y1),(1,y2),(1,y3)]
# [(2,z1),(2,z2),(2,z3)] ]

# 参考: https://zhuanlan.zhihu.com/p/352877584
相关推荐
RAG专家1 小时前
【Mixture-of-RAG】将文本和表格与大型语言模型相结合
人工智能·语言模型·rag·检索增强生成
星期天要睡觉4 小时前
自然语言处理(NLP)——自然语言处理原理、发展历程、核心技术
人工智能·自然语言处理
低音钢琴4 小时前
【人工智能系列:机器学习学习和进阶01】机器学习初学者指南:理解核心算法与应用
人工智能·算法·机器学习
飞翔的佩奇5 小时前
【完整源码+数据集+部署教程】【天线&水】舰船战舰检测与分类图像分割系统源码&数据集全套:改进yolo11-repvit
前端·python·yolo·计算机视觉·数据集·yolo11·舰船战舰检测与分类图像分割系统
大千AI助手5 小时前
Hoeffding树:数据流挖掘中的高效分类算法详解
人工智能·机器学习·分类·数据挖掘·流数据··hoeffding树
新知图书6 小时前
大模型微调定义与分类
人工智能·大模型应用开发·大模型应用
山烛6 小时前
一文读懂YOLOv4:目标检测领域的技术融合与性能突破
人工智能·yolo·目标检测·计算机视觉·yolov4
大千AI助手6 小时前
独热编码:分类数据处理的基石技术
人工智能·机器学习·分类·数据挖掘·特征工程·one-hot·独热编码
钱彬 (Qian Bin)6 小时前
项目实践4—全球证件智能识别系统(Qt客户端开发+FastAPI后端人工智能服务开发)
人工智能·qt·fastapi
钱彬 (Qian Bin)6 小时前
项目实践3—全球证件智能识别系统(Qt客户端开发+FastAPI后端人工智能服务开发)
人工智能·qt·fastapi