torch.gather以及torch,searchsorted使用示例_torch.searchsorted-优快云博客

本文链接：https://blog.youkuaiyun.com/qq_42061298/article/details/119774721

本文介绍了PyTorch中的两个重要函数：torch.gather和torch.searchsorted。torch.gather用于根据指定的索引从输入张量中收集元素，而torch.searchsorted则是在排序后的张量中找到另一个张量元素应插入的位置，返回值为一个与原张量形状相同的索引张量。这两个函数在处理张量数据时非常实用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

torch.gather使用

import torch

tensor_0 = torch.arange(3, 12).view(3, 3)
print(tensor_0)
#tensor([[ 3,  4,  5],
#        [ 6,  7,  8],
#        [ 9, 10, 11]])

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

在这里插入图片描述 **

torch.searchsorted使用

import torch.nn as nn
import torch


def main():
        x = torch.linspace(1, 11, 10)
        y = torch.linspace(2, 12, 10)
        x=x.unsqueeze(-1).expand(x.size(0),10)
        y=y.unsqueeze(-1).expand(y.size(0),100)
        print(x.shape,y.shape)
        res1 = torch.searchsorted(x, y)
        print(res1)
main()

# torch.Size([10, 10]) torch.Size([10, 100])
# tensor([[10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10],
#         ...
#         [[10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
#         10, 10, 10, 10, 10, 10, 10, 10, 10, 10]])