python - 根据 PyTorch 中的分割图在每个区域中采样

Question

我想根据段图在图像的每个区域中采样一定数量的点（就像 SLIC 算法产生的一样，它是一个与图像大小相同的图，包含从 0 到 num_segment 的整数，表示每个像素的段属于）。

目前，我编写自己的代码如下：

• 对于i范围内（0，num_segment）：
  • i使用_torch.where
  • 挑选出这些像素并形成一维张量
  • 用于torch.Upsample对i第 th 段均匀采样 n_sample 点
• 堆叠所有采样点以形成一个大的二维张量，每行代表选定的点属于一个段，它有 n_sample 行。

我不仅需要图像中每个选定点的原始值，还需要它们的索引。

我画了一张图来说明。

所以，我的问题是，在 PyTorch 中是否有本地方法来实现它？上面的代码运行有点慢，可能是for循环减慢了速度。如果可能的话，由于所有的采样过程都是相互独立的，我怎样才能加快这个过程？

一般来说，我有 ~400 个段，我想为每个段采样 20~50 个点。

标签： pythonpytorchsample