【python代码技巧2】数组索引扩增技巧

import numpy as np import torch a = np.array([x for x in range(12)]).reshape(3, 4) index = np.array([[0,1] for x in range(5)]).flatten() print "a:\n", a, a.shape print "index:\n", index, index.shape a: [[ 0 1 2 3] [ 4 5 6 7] [ 8 9 10 11]] (3, 4) index: [0 1 0 1 0 1 0 1 0 1] (10,) print a[index], a[index].shape [[0 1 2 3] [4 5 6 7] [0 1 2 3] [4 5 6 7] [0 1 2 3] [4 5 6 7] [0 1 2 3] [4 5 6 7] [0 1 2 3] [4 5 6 7]] (10, 4)

pytorch同样有这种特性

a = torch.Tensor([x for x in range(12)]).view(3, 4) index = torch.LongTensor([[0,1] for x in range(5)]).view(-1) print "a:\n", a print "index:\n", index a: 0 1 2 3 4 5 6 7 8 9 10 11 [torch.FloatTensor of size 3x4] index: 0 1 0 1 0 1 0 1 0 1 [torch.LongTensor of size 10] print a[index], a[index].size() 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 [torch.FloatTensor of size 10x4] torch.Size([10, 4])