Функция PyTorch torch.nn.Flatten

PyTorch torch.nn 参考手册Справочное руководство PyTorch torch.nn

torch.nn.FlattenЭто модуль уплощения тензора в PyTorch.

Он уплощает многомерные тензоры в одно измерение, обычно используется между свёрточными и полносвязными слоями.

Определение функции

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

Описание параметров:

  • start_dim(int): Измерение, с которого начинается уплощение. По умолчанию 1 (с сохранением размерности батча).
  • end_dim(int): Измерение, на котором заканчивается уплощение. По умолчанию -1 (до последнего измерения).

Примеры использования

Пример 1: Базовое использование

Пример

import torch
import torch.nn as nn

flatten = nn.Flatten()

# 4D-вход: (батч, каналы, высота, ширина)
x = torch.randn(4, 3, 32, 32)

output = flatten(x)

print("Форма входа:", x.shape)
print("Форма выхода:", output.shape)
print("После уплощения: 3*32*32 = 3072 измерения")

Пример 2: Сохранение размерности батча

Пример

import torch
import torch.nn as nn

# start_dim=1 сохраняет размерность батча
x = torch.randn(8, 64, 8, 8)
print("Вход:", x.shape)

# Уплощение до (8, 4096)
out1 = nn.Flatten(start_dim=1)(x)
print("Начиная с измерения 1:", out1.shape)

# Не сохранять batch
out2 = nn.Flatten(start_dim=0)(x)
print("Начиная с размерности 0:", out2.shape)

Пример 3: 3D-вход

Пример

import torch
import torch.nn as nn

# 3D-вход: (batch, seq_len, features)
x = torch.randn(4, 100, 512)

# Развернуть последовательность и признаки
flatten = nn.Flatten(start_dim=1)
output = flatten(x)

print("Форма входа:", x.shape)
print("Форма выхода:", output.shape)

Пример 4: Полный пример CNN

Пример

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("Вход:", x.shape, "-> Выход:", output.shape)

Сценарии использования

  • CNN к FC: После развертывания выхода сверточного слоя подключите к полносвязному слою
  • Преобразование размерности: Скорректируйте форму тензора

PyTorch torch.nn 参考手册Справочное руководство PyTorch torch.nn

Другие расширения