1 介绍 

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

将一个连续的维度范围扁平化为一个张量

start_dim (int)	要开始扁平化的第一个维度（默认值 = 1）
end_dim (int)	要结束扁平化的最后一个维度（默认值 = -1）

2 举例

 

input = torch.randn(32, 1, 5, 5)m = nn.Flatten()
output = m(input)
output.size()
#torch.Size([32, 25])m = nn.Flatten(0, 2)
output = m(input)
output.size()
#torch.Size([160, 5])