Pytorch实用教程：tensor.size()用法 | .squeeze()方法

Pytorch中tensor变量.size(0)

在 PyTorch 中，tensor.size(0) 是用来获取张量（Tensor）第一个维度的大小的一种方法。这里的“0”指的是第一个维度的索引，因为在 Python 和 PyTorch 中索引是从 0 开始的。换句话说，size(0) 返回的是张量在其第一个维度上的元素个数。

示例

假设我们有一个二维张量，表示一个矩阵或者一个批量的一维数据：

import torch

# 创建一个 3x4 的二维张量
x = torch.randn(3, 4)
print(x)
print(x.size(0))  # 输出张量的第一个维度的大小
1
2
3
4
5
6

如果 x 是一个 3x4 的张量，那么 x.size(0) 将会返回 3，因为它有 3 行，每一行是一个一维张量，其长度为 4。所以，这里的 3 表示的是“批量大小”或者说是这个二维张量包含的一维张量的数量。

在不同上下文中的用法

• 批量处理：在深度学习中，数据通常以批次的形式进行处理。在这种情况下，size(0) 通常用来获取批次中的样本数量。
• 多维张量：对于更高维度的张量，size(0) 依然返回第一个维度的大小，这在处理如图像数据（通常是 4D 张量，形状为 [批次大小, 通道数, 高度, 宽度]）时非常有用。

更广泛的用法

size() 方法返回一个元组，包含了张量每个维度的大小。你可以通过指定维度的索引来获取特定维度的大小，或者不传递任何参数来获取所有维度的大小：

print(x.size())  # 返回所有维度的大小
print(x.size(1))  # 返回第二个维度的大小
1
2

这种方式使得 PyTorch 在处理不同形状的张量时非常灵活和强大。

.squeeze()

在 PyTorch 中，.squeeze() 方法用于移除张量中所有维度为1的维度。当你在 .squeeze() 方法中指定一个维度参数时，它会尝试仅移除指定的维度，前提是该维度的大小确实为1。如果指定的维度不为1，则张量不会发生变化。

参数解释

• 维度参数 (dim): 当你传递一个维度给 .squeeze() 方法时，它会尝试只移除那个特定的维度。如果那个维度的大小不是1，那么原张量将保持不变。

.squeeze(-1) 的作用

当你调用 labels.squeeze(-1) 时，这意味着你想移除张量 labels 中最后一个维度（-1 指的是张量的最后一个维度），但前提是这个维度的大小为1。

• 如果 labels 的形状是 [N, M, 1]，使用 squeeze(-1) 后，它的形状将变为 [N, M]。
• 如果 labels 的最后一个维度大小不是1，比如形状是 [N, M, K] (其中 K != 1)，那么调用 squeeze(-1) 后，labels 的形状不会改变。

使用场景

这个操作在处理某些特定的数据时非常有用，例如，当你的模型输出或标签的形状为 [batch_size, num_classes, 1]，而你想将其转换为 [batch_size, num_classes] 以便计算损失函数时，这时 .squeeze(-1) 就派上了用场。

示例

让我们通过一个简单的示例来看看 .squeeze(-1) 的实际效果：

import torch

# 创建一个形状为 [3, 2, 1] 的张量
x = torch.randn(3, 2, 1)
print("Original shape:", x.shape)

# 移除最后一个维度
x_squeezed = x.squeeze(-1)
print("Shape after squeeze(-1):", x_squeezed.shape)
1
2
3
4
5
6
7
8
9

在这个示例中，x 最初的形状是 [3, 2, 1]。使用 .squeeze(-1) 后，它移除了大小为1的最后一个维度，变为了 [3, 2]。这就是 .squeeze(-1) 的作用。