pytorch的expand函数

PyTorch 中的 expand 函数用于扩展张量的形状，使其在某些维度上“看起来”像被复制了多次，但实际上它不会复制数据，从而节省内存和计算资源。扩展后的张量共享原始张量的内存空间，因此原始张量和扩展后的张量是同一个数据的视图。

以下是 torch.expand 函数的一些基本用法：

1. 扩展一维张量： 将一维张量扩展到更高维度，例如，将一维张量扩展为二维张量的第一维度。

import torch

x = torch.tensor([1, 2, 3])
y = x.expand(2, -1)  # 扩展为 2x3 的矩阵
print(y)
# 输出:
# tensor([[1, 2, 3],
#         [1, 2, 3]])

 -1 表示该维度的大小与原始张量相同。

 

2. 使用 -1 自动扩展： 使用 -1 可以自动扩展张量到与输入张量相同的大小。

x = torch.tensor([1, 2, 3])
y = x.expand(-1, 2)  # 扩展为 3x2 的矩阵
print(y)
# 输出:
# tensor([[1, 1],
#         [2, 2],
#         [3, 3]])

 

3. 扩展多维张量： 扩展多维张量时，可以指定多个维度进行扩展。

x = torch.tensor([[1, 2], [3, 4]])
y = x.expand(2, -1, -1)  # 扩展为 2x2x2 的张量
print(y)
# 输出:
# tensor([[[1, 2],
#          [1, 2]],
#         [[3, 4],
#          [3, 4]]])

 

4. 扩展与广播的区别： expand 与 broadcast 相似，但 broadcast 在进行操作时会复制数据，而 expand 不会。expand 更适用于减少内存使用。

5. 扩展与复制： 如果你需要一个实际复制了数据的新张量，可以使用 expand 后跟 clone。

x = torch.tensor([1, 2, 3])
y = x.expand(2, -1).clone()  # 扩展后复制数据
print(y)
# 输出:
# tensor([[1, 2, 3],
#         [1, 2, 3]])

 

使用 expand 时，扩展的维度大小可以是具体的数值，也可以是 -1，表示该维度的大小与原始张量相同。如果扩展的维度大小大于原始张量，PyTorch 会抛出错误。

（摘自kimi）