torch.gather、torch.squeeze、torch.unsqueeze

Torch Gather Torch Squeeze

torch.gatherバインディングtorch.squeezeを実行し、各行を抽出しますtorch.unsqueeze指定された要素を実現します（拡張可能な高次元）

import torch conf_t = torch.Tensor([[1,2,3],[2,3,4]]) print(conf_t) idx = torch.LongTensor([2, 1]) print(idx) idx2 = idx.unsqueeze(-1)# Add a dimension of -1 indicates an increase in the maximum dimension print(idx2) targets_weighted = torch.gather(conf_t, -1, idx2)In the index idx2 #, corresponding elements in aggregate conf_t print(targets_weighted) targets_weighted = targets_weighted.squeeze(-1)# 1 reduced dimensions, -1 represents the maximum dimension erasing print(targets_weighted)

出力：

tensor([[1., 2., 3.], [2., 3., 4.]]) tensor([2, 1]) tensor([[2], [1]]) tensor([[3.], [3.]]) tensor([3., 3.])