如何快速将切片从 cuda 张量 wrt 切割成另一个张量值答案

【问题标题】：How to quickly cut slices from cuda tensor wrt to another tensors values如何快速将切片从 cuda 张量 wrt 切割成另一个张量值
【发布时间】：2021-01-19 19:50:03
【问题描述】：

我有一个形状为(100, 80000, 4) 的torch cuda 张量A，还有一个形状为(1, 80000, 1) 的cuda 张量B。我想要做的是，对于第二维中的每个元素i（从0 到79999），取张量B 的值（将从0 到99 并将指出要在A 的第一个维度中取哪个值。

另一个问题是，对于B (B[0, i, 0]) 的这个元素，我想从A 中截取一个片段，即A[lower_bound:upper_bound, i:i+1, :]。

总而言之，我的张量 B 具有我想从 A 切割的切片中心的索引。我想知道，是否有办法更快地使用以下代码执行我正在执行的操作（例如，使用 cuda）

A  # tensor of shape (100, 80000, 4)
B  # tensor of shape (1, 80000, 1) which has 
k = 3  # width of the slice to take (or half of it to be exact)

A_list = []
for i in range(80000):
    lower_bound = max(0, B[0, i, 0]-k)
    upper_bound = min(100, B[0, i, 0]+k+1)
    A_mean = A[lower_bound:upper_bound, i:i+1, :].mean(0, keepdim=True)
    A_list.append(A_mean)

A = torch.cat(A_list , dim=1)

【问题讨论】：

您的代码片段中的k 是什么？
a[:, b.reshape(-1), :] 会给你ith 想要的元素。
k 是我想从 A 获取的第一维切片的宽度。在这种情况下，3.我会更新代码

标签： python python-3.x numpy pytorch

【解决方案1】：

如果 k = 1，类似的方法可以工作。（需要torch>1.7）

a = torch.rand((10, 20, 4))
b = torch.randint(10, (20, 1))
b2 = torch.cat((b-1, b, b+1), dim=1)
b2 = torch.minimum(9*torch.ones_like(b2), b2)
b2 = torch.maximum(0*torch.ones_like(b2), b2)
a[:, b2, :]

然后重新整形以获得正确的尺寸

【讨论】：