深圳幻海软件技术有限公司 欢迎您!

nn.Flatten()函数详解及示例

2023-02-28

torch.nn.Flatten(start_dim=1, end_dim=- 1)作用:将连续的维度范围展平为张量。经常在nn.Sequential()中出现,一般写在某个神经网络模型之后,用于对神经网络模型的输出进行处理,得到tensor类型的数据。 有俩个参数,s

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

作用:将连续的维度范围展平为张量。 经常在nn.Sequential()中出现,一般写在某个神经网络模型之后,用于对神经网络模型的输出进行处理,得到tensor类型的数据。

 有俩个参数,start_dim和end_dim,分别表示开始的维度和终止的维度,默认值分别是1和-1,其中1表示第一维度,-1表示最后的维度。结合起来看意思就是从第一维度到最后一个维度全部给展平为张量。(注意:数据的维度是从0开始的,也就是存在第0维度,第一维度并不是真正意义上的第一个)

同理,如果我这么写:

self.flat = nn.Flatten(start_dim=2, end_dim=3)

那么意思就是从第二维度开始,到第三维度全部给展平,也就是将2、3两个维度展平。

官网给出的示例:

  1. input = torch.randn(32, 1, 5, 5)
  2. # With default parameters
  3. m = nn.Flatten()
  4. output = m(input)
  5. output.size()
  6. #torch.Size([32, 25])
  7. # With non-default parameters
  8. m = nn.Flatten(0, 2)
  9. output = m(input)
  10. output.size()
  11. #torch.Size([160, 5])

#开头的代码是注释

整段代码的意思是:给定一个维度为(32,1,5,5)的随机数据。

1.先使用一次nn.Flatten(),使用默认参数:

m = nn.Flatten()

也就是说从第一维度展平到最后一个维度,数据的维度是从0开始的,第一维度实际上是数据的第二个位置代表的维度,也就是样例中的1。

因此进行展平后的结果也就是[32,1×5×5]➡[32,25]

2.接着再使用一次指定参数的nn.Flatten(),即

m = nn.Flatten(0, 2)

也就是说从第0维度展平到第2维度,0~2,对应的也就是前三个维度。

因此结果就是[32×1×5,5]➡[160,5]

因此进行展平后的结果也就是[32,1*5*5]➡[32,25]

文章知识点与官方知识档案匹配,可进一步学习相关知识
Python入门技能树人工智能深度学习237982 人正在系统学习中