PyTorch torch.nn.Flatten Function

PyTorch torch.nn 参考手册PyTorch torch.nn Reference Manual

torch.nn.FlattenIt is a tensor flattening module in PyTorch.

It flattens multi-dimensional tensors into one dimension, commonly used between convolutional layers and fully connected layers.

Function Definition

torch.nn.Flatten(start_dim=1, end_dim=-1)

Parameter Description:

  • start_dim(int): The dimension to start flattening. Default is 1 (preserving the batch dimension).
  • end_dim(int): The dimension to end flattening. Default is -1 (up to the last dimension).

Usage Examples

Example 1: Basic Usage

Example

import torch
import torch.nn as nn

flatten = nn.Flatten()

# 4D input: (batch, channels, height, width)
x = torch.randn(4, 3, 32, 32)

output = flatten(x)

print("Input shape:", x.shape)
print("Output shape:", output.shape)
print("After flattening: 3*32*32 = 3072 dimensions")

Example 2: Preserving the Batch Dimension

Example

import torch
import torch.nn as nn

# start_dim=1 preserves the batch dimension
x = torch.randn(8, 64, 8, 8)
print("Input:", x.shape)

# Flatten to (8, 4096)
out1 = nn.Flatten(start_dim=1)(x)
print("Starting from dimension 1:", out1.shape)

# Do not preserve batch
out2 = nn.Flatten(start_dim=0)(x)
print("Starting from dimension 0:", out2.shape)

Example 3: 3D Input

Example

import torch
import torch.nn as nn

# 3D input: (batch, seq_len, features)
x = torch.randn(4, 100, 512)

# Flatten the sequence and features
flatten = nn.Flatten(start_dim=1)
output = flatten(x)

print("Input shape:", x.shape)
print("Output shape:", output.shape)

Example 4: Complete CNN Example

Example

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Conv2d(3, 32, 3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Conv2d(32, 64, 3, padding=1),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d(1),
    nn.Flatten(),
    nn.Linear(64, 10)
)

x = torch.randn(4, 3, 32, 32)
output = model(x)

print("Input:", x.shape, "-> Output:", output.shape)

Use Cases

  • CNN to FC: After flattening the convolutional layer output, connect to the fully connected layer
  • Dimension Transformation: Adjust tensor shape

PyTorch torch.nn 参考手册PyTorch torch.nn Reference Manual

Other Extensions