Python torch.unsqueeze的使用

什么是torch.unsqueeze

torch.unsqueeze函数是PyTorch中的一个方法，用于在指定维度上增加一个维度。它可以将一个大小为 n 的Tensor转换为一个大小为 1×n 的二维Tensor。

torch.unsqueeze的语法

torch.unsqueeze(input, dim)

input：输入的Tensor。
dim：指定要添加的维度的索引。

torch.unsqueeze的使用示例

下面是一个使用torch.unsqueeze的示例，以帮助我们更好地理解它的使用方法。

import torch

# 定义一个二维Tensor
tensor = torch.tensor([[1, 2], [3, 4], [5, 6]])

print('原始Tensor大小：', tensor.size())

# 使用torch.unsqueeze在dim=1上增加一个维度
unsqueeze_tensor = torch.unsqueeze(tensor, 1)

print('添加维度后的Tensor大小：', unsqueeze_tensor.size())
print('添加维度后的Tensor内容：\n', unsqueeze_tensor)

运行上述代码，输出结果如下：

原始Tensor大小： torch.Size([3, 2])
添加维度后的Tensor大小： torch.Size([3, 1, 2])
添加维度后的Tensor内容：
tensor([[[1, 2]],

        [[3, 4]],

        [[5, 6]]])

从输出结果可以看出，原始Tensor是一个大小为3×2的二维Tensor。使用torch.unsqueeze函数，在dim=1的位置添加了一个维度，结果得到了一个大小为3×1×2的三维Tensor。注意，在输出结果中的每个Tensor都用三个方括号表示，这表示Tensor的维度。每个方括号对应一个维度。

torch.unsqueeze的应用场景

torch.unsqueeze的主要应用场景之一是在深度学习模型中数据的预处理过程中。在某些情况下，我们需要将一维或二维数据转换为更高维度的形式，以便与模型的输入要求相匹配。

下面是一个示例，在卷积神经网络（CNN）中使用torch.unsqueeze将一个一维向量转换为三维输入。

import torch
import torch.nn as nn

# 定义一个简单的卷积神经网络
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=3)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=3)
        self.fc = nn.Linear(20, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# 创建一个模型实例
model = SimpleCNN()

# 定义一个一维输入向量
input = torch.tensor([1, 2, 3, 4, 5])

print('原始输入向量大小：', input.size())

# 使用torch.unsqueeze在dim=0上增加一个维度
unsqueeze_input = torch.unsqueeze(input, 0)

print('添加维度后的输入向量大小：', unsqueeze_input.size())
print('添加维度后的输入向量内容：\n', unsqueeze_input)

# 前向传播计算输出
output = model(unsqueeze_input)

print('模型输出大小：', output.size())
print('模型输出结果：', output)

运行上述代码，输出结果如下：

原始输入向量大小： torch.Size([5])
添加维度后的输入向量大小： torch.Size([1, 5])
添加维度后的输入向量内容：
tensor([[1, 2, 3, 4, 5]])
模型输出大小： torch.Size([1, 10])
模型输出结果： tensor([[0.0277, 0.0105, 0.1271, 0.0153, 0.0397, 0.1043, 0.1238, 0.0719, 0.0669,
         0.0439]], grad_fn=<AddmmBackward>)

从输出结果可以看出，原始的一维输入向量大小为5。使用torch.unsqueeze函数，在dim=0的位置添加了一个维度，结果得到了一个大小为1×5的二维输入。经过模型的前向传播计算后，得到了一个大小为1×10的输出。