避免在 PyTorch 的批量张量中选择维度的循环答案

【问题标题】：Avoid Loop for Selecting Dimensions in a Batch Tensor in PyTorch避免在 PyTorch 的批量张量中选择维度的循环
【发布时间】：2021-09-16 16:49:46
【问题描述】：

我有一个批量张量和另一个具有维度索引的张量，可以从批量张量中选择。目前，我正在循环批处理张量，如下代码sn-p所示：

import torch

# create tensors to represent our data in torch format
batch_size = 8
batch_data = torch.rand(batch_size, 3, 240, 320)

# notice that channels_id has 8 elements, i.e., = batch_size
channels_id = torch.tensor([2, 0, 2, 1, 0, 2, 1, 0])

这就是我在 for 循环中选择尺寸然后堆叠以转换单个张量的方式：

batch_out = torch.stack([batch_i[channel_i] for batch_i, channel_i in zip(batch_data, channels_id)])
batch_out.size()  # prints torch.Size([8, 240, 320])

效果很好。但是，有没有更好的 PyTorch 方法来实现同样的效果？

【问题讨论】：

矢量化函数在 numpy 中处理得更好。在您的情况下，可以在 numpy 中执行操作，然后转换为 Torch 张量。但缺点是你不能使用 gpu 功能。
看torch.gather

标签： python indexing pytorch tensor

【解决方案1】：

根据@Shai 的提示，我可以使用torch.gather 函数使其工作。以下是完整代码：

import torch

# create tensors to represent our data in torch format
batch_size = 8
batch_data = torch.rand(batch_size, 3, 240, 320)

# notice that channels_id has 8 elements, i.e., batch_size
channels_id = torch.tensor([2, 0, 2, 1, 0, 2, 1, 0])

# resizing channels_id to (8 , 1, 240, 320)
channels_id = channels_id.view(-1, 1, 1, 1).repeat((1, 1) + batch_data.size()[-2:])

batch_out = torch.gather(batch_data, 1, channels_id).squeeze()
batch_out.size()  # prints torch.Size([8, 240, 320])

【讨论】：