python切片学习：a[b，c]a，b，c都是tensor，用tensor来切片tensor

今天在yolov5中看到不少这个操作，但是搜了一下csdn很少有人详细解释这种，于是自己测试了一下。本人小白，如果有不对的地方敬请指正。代码是用了文章末尾链接里的代码修改来的。那个文章讲的很好，大家也可以看一下

问题：
yolov5代码：在loss.py里面
pi是一个tensor [3,16,128,128,200]
a，b，gi，gj也分别是一个tensor 都是一行1408列 [1280,]

求: ps = pi[b, a, gj, gi] 
 
输出是[1408，200]的tensor

其中[b, a, gj, gi]都是一维相同长度的tensor。

测试学习如下：

import torch
import numpy as np

y = np.array([1,2])
z = np.array([1,2])
a = torch.tensor(y, dtype=torch.long)
b = torch.tensor(z, dtype=torch.long)
x = torch.rand(5,3,2)
print(x)
print(y)
print(z)
print(x[a,b])   

#x = tensor([[[0.3957, 0.2314],
#          [0.9333, 0.7191],
#          [0.4622, 0.9004]],
# 
#         [[0.2074, 0.0972],
#          [0.9954, 0.5551],
#          [0.4886, 0.0336]],
# 
#         [[0.5855, 0.2486],
#          [0.3410, 0.7941],
#          [0.0258, 0.4987]],
# 
#         [[0.6768, 0.6552],
#          [0.2251, 0.9043],
#          [0.2961, 0.4868]],
# 
#         [[0.7500, 0.6140],
#          [0.9636, 0.8335],
#          [0.1801, 0.4498]]])
# y = [1 2]
# z = [1 2]
# x[a,b] = tensor([[0.9954, 0.5551],
#                 [0.0258, 0.4987]])

可以把[a,b]想象成一个矩阵第一行是（1，1）对应第一个维度的第二堆和第二个维度的第二堆也就是[0.9954,0.5551]

矩阵的第二行（2，2）对应第一个维度的第三堆数和第二个维度的即第三堆数中的第三行即[0.0258,0.4987]

结果：把[a,b,gj,gi]想象成一个1408×4矩阵每行有四个数每列有1408个数

那四个数分别对应pi的前四个维度然后就可以取出列为200的四个数代表的行

一共可以取出来1408个行

输出的维度即为[1408,200]

参考：

https://blog.csdn.net/weixin_50764885/article/details/128311289

https://blog.csdn.net/qq_38929105/article/details/117994129