PyTorch torch.nn.Flatten 関数

PyTorch torch.nn リファレンスマニュアル

torch.nn.Flattenこれは PyTorch のテンソル平坦化モジュールです。

これは多次元テンソルを1次元に平坦化し、畳み込み層と全結合層の間の接続によく使用されます。

関数定義

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

パラメータの説明:

  • start_dim(int): 平坦化を開始する次元。デフォルトは 1（batch 次元を保持）。
  • end_dim(int): 平坦化を終了する次元。デフォルトは -1（最後の次元まで）。

使用例

例 1: 基本的な使い方

サンプル

import torch
import torch.nn as nn

flatten = nn.Flatten()

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

output = flatten(x)

print("入力形状:", x.shape)
print("出力形状:", output.shape)
print("平坦化後: 3*32*32 = 3072 次元")

例 2: batch 次元を保持

サンプル

import torch
import torch.nn as nn

# start_dim=1 は batch 次元を保持
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 リファレンスマニュアル

その他の拡張