pytorch tensor.expand函数介绍

在 PyTorch 中，tensor.expand()是一个用于扩展张量维度的函数。

一、函数作用

它允许你在不复制数据的情况下，将张量的形状扩展到指定的维度大小。这对于需要在特定维度上重复数据的操作非常有用，例如在进行广播操作时调整张量的形状。

二、函数语法

tensor.expand(*sizes)

其中，*sizes是一个可变参数，表示要扩展到的目标形状。可以传入整数或整数序列来指定每个维度的大小。

三、使用示例

示例1:

import torch

# 创建一个形状为 (2,) 的一维张量
tensor = torch.tensor([1, 2])

# 将张量扩展为形状为 (2, 3) 的二维张量
expanded_tensor = tensor.expand(2, 3)
print(expanded_tensor)

输出：

tensor([[1, 2, 1],
        [1, 2, 1]])

示例2:

import torch

# 创建一个三维张量
tensor = torch.randn(2, 1, 3)
print("原始张量形状：", tensor.shape)

# 扩展张量的维度
expanded_tensor = tensor.expand(2, 2, 3)
print("扩展后张量形状：", expanded_tensor.shape)