torch.repeat()、torch.repeat_interleave()和torch.expand()

1. torch.repeat_interleave(input, repeats, dim=None) 该函数用来复制张量元素

• 参数
  – input： 输入张量
  – repeats： 对张量的每个元素复制的次数，通过广播机制实现，可以是int类型，也可以是张量类型。指定每行的复制次数。
  – dim：指定对哪个轴上的元素进行复制，默认将输入数组展开复制并返回一个展开后的输出数组

1.1 不指定轴

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
a.repeat_interleave(2)
# tensor([1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6])

1.2 指定轴

1.2.1 按行

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
a.repeat_interleave(2, 1)
# tensor([[1, 1, 2, 2, 3, 3],
#        [4, 4, 5, 5, 6, 6]])

1.2.2 按列

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
a.repeat_interleave(2, 0)
tensor([[1, 2, 3],
        [1, 2, 3],
        [4, 5, 6],
        [4, 5, 6]])

1.3 使用张量指定行的复制次数

aa = torch.tensor([[1, 2], [3, 4], [5, 6]])
torch.repeat_interleave(aa, torch.tensor([3, 2, 2]), dim=0)
tensor([[1, 2],
        [1, 2],
        [1, 2],
        [3, 4],
        [3, 4],
        [5, 6],
        [5, 6]])

2 torch.repeat(*sizes) 整行/整列复制

• 例子一

bb = torch.tensor([1, 2, 3])
bb.repeat(4, 2)
tensor([[1, 2, 3, 1, 2, 3],
        [1, 2, 3, 1, 2, 3],
        [1, 2, 3, 1, 2, 3],
        [1, 2, 3, 1, 2, 3]])

• 例子二

cc = bb.repeat(4, 2, 1)
cc.shape  # (4, 2, 3)
tensor([[[1, 2, 3],
         [1, 2, 3]],

        [[1, 2, 3],
         [1, 2, 3]],

        [[1, 2, 3],
         [1, 2, 3]],

        [[1, 2, 3],
         [1, 2, 3]]])
         

3. torch.expand(*sizes)

返回原张量扩容后的一个新的视图

• 例子一

x = torch.tensor([[1], [2], [3]])
x.size()  # (3, 1)
b = x.expand(3, 4)  # 等价于x.expand(-1, 4), -1 代表不变
tensor([[1, 1, 1, 1],
        [2, 2, 2, 2],
        [3, 3, 3, 3]])