Fonction PyTorch torch.nn.Flatten

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

torch.nn.FlattenC'est un module d'aplatissement de tenseur dans PyTorch.

Il aplatit les tenseurs multidimensionnels en une dimension, couramment utilisé entre les couches convolutives et les couches entièrement connectées.

Définition de fonction

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

Description des paramètres :

  • start_dim(int) : La dimension de début de l'aplatissement. La valeur par défaut est 1 (en préservant la dimension de lot).
  • end_dim(int) : La dimension de fin de l'aplatissement. La valeur par défaut est -1 (jusqu'à la dernière dimension).

Exemples d'utilisation

Exemple 1 : Utilisation de base

Exemple

import torch
import torch.nn as nn

flatten = nn.Flatten()

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

output = flatten(x)

print("Forme d'entrée :", x.shape)
print("Forme de sortie :", output.shape)
print("Après aplatissement : 3*32*32 = 3072 dimensions")

Exemple 2 : Conservation de la dimension de lot

Exemple

import torch
import torch.nn as nn

# start_dim=1 préserve la dimension de lot
x = torch.randn(8, 64, 8, 8)
print("Entrée :", x.shape)

# Aplatir vers (8, 4096)
out1 = nn.Flatten(start_dim=1)(x)
print("À partir de la dimension 1 :", out1.shape)

# Ne pas préserver le batch
out2 = nn.Flatten(start_dim=0)(x)
print("À partir de la dimension 0 :", out2.shape)

Exemple 3 : Entrée 3D

Exemple

import torch
import torch.nn as nn

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

# Aplatir la séquence et les caractéristiques
flatten = nn.Flatten(start_dim=1)
output = flatten(x)

print("Forme d'entrée :", x.shape)
print("Forme de sortie :", output.shape)

Exemple 4 : Exemple complet de CNN

Exemple

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("Entrée :", x.shape, "-> Sortie :", output.shape)

Cas d'utilisation

  • CNN vers FC: Après avoir aplati la sortie de la couche convolutionnelle, connecter à la couche entièrement connectée
  • Transformation de dimension: Ajuster la forme du tenseur

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

Autres extensions