Pytorch的gather和scatter
最近在看pytorch的gather与scatter函数，现在详细记录一下
1、Gather
gather是根据索引取数据，下图可以表示gather，具体见[gather]的介绍(https://stackoverflow.com/a/54706716)
但是要注意的是，dim为0和1时，index是有区别的，要转置一下
index = torch.as_tensor([[0,1,2],[1,2,0]])
src = torch.arange(9).reshape(3,3)
torch.gather(src,0,index)
torch.gather(src,1,index.T) #dim 为1时，index要转置
2、Scatter
scatter是将数据根据索引回填到新的矩阵里面，这个适合做onehot矩阵
1)对于2D转3D
如下图，进行回填，参考知乎
a = torch.rand(2, 5)
print(a)
b = torch.zeros(3, 5).scatter_(0, torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]), a)
print(b)
制作one-hot的代码，参考PyTorch笔记之 scatter() 函数
class_num = 10
batch_size = 4
label = torch.LongTensor(batch_size, 1).random_() % class_num
#tensor([[6],
# [0],
# [3],
# [2]])
torch.zeros(batch_size, class_num).scatter_(1, label, 1)
#tensor([[0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
# [1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
# [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
# [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.]])
1)对于3D转4D的
这个对于图像分割的onehot制作比较多，这时，每个类别在一个面上设置为0或1，具体参考PyTorch One-Hot Labels
def make_one_hot(labels, C=2):
'''
Converts an integer label torch.autograd.Variable to a one-hot Variable.
Parameters
----------
labels : torch.autograd.Variable of torch.cuda.LongTensor
N x 1 x H x W, where N is batch size.
Each value is an integer representing correct classification.
C : integer.
number of classes in labels.
Returns
-------
target : torch.autograd.Variable of torch.cuda.FloatTensor
N x C x H x W, where C is class number. One-hot encoded.
'''
one_hot = torch.cuda.FloatTensor(labels.size(0), C, labels.size(2), labels.size(3)).zero_()
target = one_hot.scatter_(1, labels.data, 1)
target = Variable(target)
return target
具体结果为
>> labels = torch.LongTensor(4,4) % 3
2 1 0 0
1 0 0 0
2 0 0 1
2 0 0 1
[torch.LongTensor of size 4x4]
>> make_one_hot(labels)
(0 ,0 ,.,.) =
0 0 1 1
0 1 1 1
0 1 1 0
0 1 1 0
(0 ,1 ,.,.) =
0 0 0 0
1 0 0 0
0 0 0 1
0 0 0 1
(0 ,2 ,.,.) =
1 1 0 0
0 0 0 0
1 0 0 0
1 0 0 0
[torch.LongTensor of size 1x3x4x4]
这个由于是三维图像，所以要立体的看，竖着从上往下看，看这个位置中，1放在第几层
