【PyTorch】torch.flatten()函数：将多维张量展平为一个一维张量

在 PyTorch 中，torch.flatten() 是一个非常实用的函数，用于将张量展平（flatten）为一维张量。展平的操作将多维张量转换为一个一维张量，但保留所有的元素。它通常用于神经网络模型中，将多维输入展平为一维，以便送入全连接层（Fully Connected Layer）。

语法

torch.flatten(input, start_dim=0, end_dim=-1)

参数：

• input：输入张量，形状可以是任意维度。
• start_dim（可选）：开始展平的维度。默认值是 0，表示从第一个维度开始展平。
• end_dim（可选）：结束展平的维度。默认值是 -1，表示直到最后一个维度进行展平。

返回：

返回展平后的张量。

例子

1. 基本的展平操作

import torch

# 创建一个 2x3x4 的张量
x = torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
                  [[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]]])

# 展平整个张量
flat_x = torch.flatten(x)
print(flat_x)

输出：

tensor([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24])

2. 从指定维度开始展平

你可以指定从某个维度开始展平。例如，指定从第二维开始展平：

import torch

x = torch.tensor([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]])

# 从维度 1 开始展平
flat_x = torch.flatten(x, start_dim=1)
print(flat_x)

输出：

tensor([ [ 1,  2,  3,  4,  5,  6],
         [ 7,  8,  9, 10, 11, 12]])

3. 展平多个维度

如果你想展平从某个维度到另一个维度之间的部分，可以使用 start_dim 和 end_dim 参数。比如，展平第二维到第三维：

import torch

x = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])

# 展平从维度 1 到 2（包括第2维和第3维）
flat_x = torch.flatten(x, start_dim=1, end_dim=2)
print(flat_x)

输出：

tensor([[1, 2, 3, 4],
        [5, 6, 7, 8]])

4. 常见用法

在神经网络中，torch.flatten() 常用于将卷积层的输出展平，以便送入全连接层。假设卷积层的输出是一个多维张量（例如 batch_size x channels x height x width），你可以使用 torch.flatten() 将其展平为一维，然后送入全连接层。

import torch
import torch.nn as nn

# 假设一个卷积层的输出张量，形状为 (batch_size, channels, height, width)
x = torch.randn(32, 64, 8, 8)  # 假设 batch_size=32, channels=64, height=8, width=8

# 使用 flatten 展平为一维
flattened_x = torch.flatten(x, start_dim=1)  # 忽略 batch 维度，从第二维开始展平

print(flattened_x.shape)  # 结果形状: (32, 4096)  64*8*8 = 4096

5. 与 view() 函数的区别

torch.flatten() 与 view() 都可以用来展平张量，但 flatten() 是一个专门的函数，通常更容易理解和使用。如果只是简单的展平整个张量，flatten() 比 view() 更直观。

view() 方法可以将张量的形状更改为给定的形状，但需要确保给定的新形状与原始张量的元素数量匹配。

# 使用 view 展平张量
flat_x = x.view(-1)  # -1 表示自动推导出维度

总结：

• torch.flatten() 是一个专门用于展平张量的函数。
• 它非常适合用来将卷积层的输出展平送入全连接层。
• flatten() 与 view() 都可以用来调整形状，但 flatten() 更简洁，且通常不需要担心元素数量是否匹配。