Función PyTorch torch.nn.Flatten

PyTorch torch.nn 参考手册Manual de referencia de PyTorch torch.nn

torch.nn.FlattenEs un módulo de aplanamiento de tensores en PyTorch.

Aplana tensores multidimensionales a una dimensión, utilizado comúnmente para la conexión entre capas convolucionales y capas completamente conectadas.

Definición de la función

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

Descripción de parámetros:

  • start_dim(int): Dimensión donde comienza el aplanamiento. Por defecto es 1 (conserva la dimensión batch).
  • end_dim(int): Dimensión donde termina el aplanamiento. Por defecto es -1 (hasta la última dimensión).

Ejemplo de uso

Ejemplo 1: Uso básico

Ejemplo

import torch
import torch.nn as nn

flatten = nn.Flatten()

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

output = flatten(x)

print("Forma de entrada:", x.shape)
print("Forma de salida:", output.shape)
print("Después de aplanar: 3*32*32 = 3072 dimensiones")

Ejemplo 2: Conservar la dimensión batch

Ejemplo

import torch
import torch.nn as nn

# start_dim=1 conserva la dimensión batch
x = torch.randn(8, 64, 8, 8)
print("Entrada:", x.shape)

# Aplanar a (8, 4096)
out1 = nn.Flatten(start_dim=1)(x)
print("Desde la dimensión 1:", out1.shape)

# Sin conservar batch
out2 = nn.Flatten(start_dim=0)(x)
print("Desde la dimensión 0:", out2.shape)

Ejemplo 3: Entrada 3D

Ejemplo

import torch
import torch.nn as nn

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

# Aplanar secuencia y características
flatten = nn.Flatten(start_dim=1)
output = flatten(x)

print("Forma de entrada:", x.shape)
print("Forma de salida:", output.shape)

Ejemplo 4: Ejemplo completo de CNN

Ejemplo

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("Entrada:", x.shape, "-> Salida:", output.shape)

Escenarios de uso

  • CNN a FC: La salida de la capa convolucional se aplana y se conecta a la capa completamente conectada
  • Transformación de dimensiones: Ajustar la forma del tensor

PyTorch torch.nn 参考手册Manual de referencia de PyTorch torch.nn

Otras extensiones