注意默认的start_dim从1开始
https://pytorch.org/docs/stable/generated/torch.nn.Flatten.html?highlight=flatten#torch.nn.Flatten
示例1
卷积公式
o u t s i z e = i n p u t s i z e − k s i z e + 2 ∗ p a d d i n g s t r i d e + 1 outsize=\frac{inputsize - ksize + 2*padding}{stride} + 1 outsize=strideinputsize−ksize+2∗padding+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])