pytorch——unsqueeze与expand


torch中的unsqueeze()函数来增加一个维度，expand()函数以行或列来广播。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# -*- encoding: utf-8 -*-
import torch

# 需求是对一个batch_size=2, seq_len=3的两个序列进行mask的扩展，
# 扩展为[batch_size, seq_len, 4, seq_len]
tokens = torch.tensor([[1,2, 3],[2,1,0]])
mask = tokens!=0
print(mask)
print(mask.shape)

print(mask.unsqueeze(2).shape)
print(mask.unsqueeze(2))
print(mask.unsqueeze(1).shape)
print(mask.unsqueeze(1))

multi = mask.unsqueeze(2)*mask.unsqueeze(1)
print('multi shape:',multi.shape) # [batch_size, seq, seq]
print(multi)

select = multi.unsqueeze(2)
print(select.shape) # batch, seq, 1, seq
print(select)
print(select.expand(-1,-1, 4, -1)) # expand的作用是把某个维度上为1的扩展为指定的个数
  • expand()在行或列上的扩展
1
2
3
4
5
6
7
8
9
10
b shape: torch.Size([3, 1])
bb shape: torch.Size([3, 3])
tensor([[1, 1, 1],
        [2, 2, 2],
        [3, 3, 3]])
c shape: torch.Size([1, 3])
cc shape: torch.Size([3, 3])
tensor([[1, 2, 3],
        [1, 2, 3],
        [1, 2, 3]])

result:

1
2
3
4
5
6
7
8
9
10
b shape: torch.Size([3, 1])
bb shape: torch.Size([3, 3])
tensor([[1, 1, 1],
        [2, 2, 2],
        [3, 3, 3]])
c shape: torch.Size([1, 3])
cc shape: torch.Size([3, 3])
tensor([[1, 2, 3],
        [1, 2, 3],
        [1, 2, 3]])