nn.Flatten()函数

最新推荐文章于 2025-04-27 14:49:16 发布

和猫君共建乌托邦~

最新推荐文章于 2025-04-27 14:49:16 发布

阅读量1.2w

点赞数 8

文章标签： pytorch 深度学习神经网络

本文链接：https://blog.youkuaiyun.com/jiacong_wang/article/details/122941990

版权

本文介绍了PyTorch中Flatten层的作用，它用于将多维输入展平为一维。通过示例展示了在卷积神经网络后使用Flatten层如何改变输出尺寸。示例1中，默认的start_dim=1，输出大小为[32, 288]；示例2中，指定start_dim=0，输出大小变为[9216]。这两个例子突显了start_dim参数的影响。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

注意默认的start_dim从1开始
https://pytorch.org/docs/stable/generated/torch.nn.Flatten.html?highlight=flatten#torch.nn.Flatten

示例1

卷积公式

$outsize=\frac{inputsize - ksize + 2*padding}{stride} + 1$

import torch
import torch.nn as nn
input = torch.randn(32, 1, 5, 5)
m = nn.Sequential(
    nn.Conv2d(1, 32, 5, 1, 1),  # 通过卷积，得到torch.size([32, 32, 3, 3]
    nn.Flatten())

output = m(input)
print(output.size())

>> torch.Size([32, 288])

示例2

import torch
import torch.nn as nn
input = torch.randn(32, 1, 5, 5)
m = nn.Sequential(
    nn.Conv2d(1, 32, 5, 1, 1),  # 通过卷积，得到torch.size([32, 32, 3, 3]
    nn.Flatten(start_dim=0))

output = m(input)
print(output.size())

>>torch.Size([9216])