PyTorch Tensor 变换

1
import torch
2

3
tensor = torch.LongTensor([[[1,2,3],[4,5,6]], [[7,8,9],[10,11,12]]])
4
print(tensor)
5
# tensor([[[ 1,  2,  3],
6
#          [ 4,  5,  6]],
7
#
8
#         [[ 7,  8,  9],
9
#          [10, 11, 12]]])
10

11
# 1. Shape & dtype basics
12
print(tensor.shape)
13
# torch.Size([2, 2, 3])
14

15

16
# 2. Reshape family (no math change)
17
# reshape 这里有个使用点是，-1 可以自动计算，比如这里是 2 * 6 = 12；
18
print(tensor.reshape(-1, 6))
19
# tensor([[ 1,  2,  3,  4,  5,  6],
20
#         [ 7,  8,  9, 10, 11, 12]])
21
print(tensor.reshape(4, 3))
22
# tensor([[ 1,  2,  3],
23
#         [ 4,  5,  6],
24
#         [ 7,  8,  9],
25
#         [10, 11, 12]])
26

27
# 3. Reorder dimensions
28
print(tensor.permute(0, 2, 1).shape)
29
# torch.Size([2, 3, 2])
30
print(tensor.permute(0, 2, 1))
31
# tensor([[[ 1,  4],
32
#          [ 2,  5],
33
#          [ 3,  6]],
34
#
35
#         [[ 7, 10],
36
#          [ 8, 11],
37
#          [ 9, 12]]])
38

39
# 4. Add/remove size-1 dims
40
print(tensor.unsqueeze(0))
41
# tensor([[[[ 1,  2,  3],
42
#           [ 4,  5,  6]],
43
#
44
#          [[ 7,  8,  9],
45
#           [10, 11, 12]]]])
46
print(tensor.squeeze(0))
47
# tensor([[[ 1,  2,  3],
48
#          [ 4,  5,  6]],
49
#
50
#         [[ 7,  8,  9],
51
#          [10, 11, 12]]])
52

53

54
# 5. Broadcasting
55
tensor = torch.FloatTensor([[1,2,3],[4,5,6]])
56
mean = tensor.mean(dim=1)
57
print(mean)
58
# tensor([2.0000, 5.0000])
59
mean = mean.unsqueeze(1).expand_as(tensor)
60
print(mean)
61
# tensor([[2., 2., 2.],
62
#         [5., 5., 5.]])
63

64
# 6. Indexing & slicing
65
tensor = torch.randn(2,3,4,5)
66
print(tensor.shape)
67
# torch.Size([2, 3, 4, 5])
68
y = tensor[:, 0]
69
print(y.shape)
70
# torch.Size([2, 4, 5])
71
y = tensor[..., -1]
72
print(y.shape)
73
# torch.Size([2, 3, 4])
74

75
# mask
76
tensor = torch.LongTensor([1,2,3,4,5,6])
77
y = torch.zeros_like(tensor)
78
y = torch.where(tensor > 3, tensor, y)
79
print(y)
80
# tensor([0, 0, 0, 4, 5, 6])
81

82
# 7. Combine / split
83
# cat
84
a = torch.randn(2,3)
85
b = torch.randn(2,3)
86
c = torch.cat([a,b], dim=0)
87
print(c.shape)
88
# torch.Size([4, 3])
89
c = torch.cat([a,b], dim=1)
90
print(c.shape)
91
# torch.Size([2, 6])
92

93
# stack (create a new dimension)
94
c = torch.stack([a,b], dim=0)
95
print(c.shape)
96
# torch.Size([2, 2, 3])
97
c = torch.stack([a,b,a], dim=1)
98
print(c.shape)
99
# torch.Size([2, 3, 3])
100
c = torch.stack([a,b], dim=2)
101
print(c.shape)
102
# torch.Size([2, 3, 2])
103

104
# expand vs repeat
105
# expand does not copy data, it only creates a new view of the same data
106
# repeat copies data
107
a = torch.randn(2,3)
108
b = torch.randn(2,7)
109
c = a.expand_as(b)
110
print(c.shape)
111
# torch.Size([2, 7])
112
c = a.repeat_interleave(2, dim=0)
113
print(c)
114
# tensor([[1, 2, 3],
115
#         [1, 2, 3],
116
#         [4, 5, 6],
117
#         [4, 5, 6]])
118

119

120
# 8. Memory & performance (important!)
121
tensor = torch.LongTensor([[1,2,3],[4,5,6]])
122

123
# How many memory steps to move when increasing a dimension index by 1
124
print(tensor.stride())
125
# (3, 1)
126

127
# view does NOT move data in memory, it only reinterprets the existing memory layout
128
y = tensor.view(-1)
129
print(y)
130
print(y.stride())
131

132
# permute changes the order of dimensions and the stride
133
y = tensor.permute(1, 0)
134
print(y)
135
print(y.stride())
136
# (1, 3)
137

138
# after permute, the order of elements is changed in memory, we can't use view to reshape it
139
# do not execute y.view(-1), use y.contiguous() to get a contiguous tensor first
140
y = tensor.permute(1, 0).contiguous().view(-1)
141
print(y)
142
# tensor([1, 2, 3, 4, 5, 6])