torch 操作函数

torch.multinomial

torch.multinomial用于从多项分布中采样。给定一个包含概率的张量，该函数会根据这些概率返回采样的索引。适用于需要根据概率分布进行随机选择的场景，例如在强化学习中根据策略选择动作。

输入是一个概率分布张量，每个元素代表对应索引的概率。函数返回一个张量，包含根据这些概率采样的索引。参数num_samples指定采样的数量，replacement决定是否允许重复采样。

importtorch probs=torch.tensor([0.1,0.2,0.7])samples=torch.multinomial(probs,num_samples=5,replacement=True)print(samples)

torch.gather

torch.gather用于根据索引从输入张量中收集元素。适用于需要根据动态索引从张量中提取数据的场景，例如在序列模型中根据注意力权重选择特征。

输入包括源张量和索引张量，索引张量的形状决定了输出张量的形状。dim参数指定沿哪个维度进行收集操作。索引张量的每个元素对应源张量中某个位置的索引。

importtorchinput=torch.tensor([[1,2],[3,4]])index=torch.tensor([[0,0],[1,0]])output=torch.gather(input,dim=1,index=index)print(output)