pytorch中nn.Flatten()函數(shù)詳解及示例

更新時(shí)間：2023年01月06日 16:02:12 作者：淺摯灬半離兮

nn.Flatten是一個(gè)類,而torch.flatten()則是一個(gè)函數(shù),下面這篇文章主要給大家介紹了關(guān)于pytorch中nn.Flatten()函數(shù)詳解及示例的相關(guān)資料,需要的朋友可以參考下

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

作用：將連續(xù)的維度范圍展平為張量。經(jīng)常在nn.Sequential()中出現(xiàn)，一般寫在某個(gè)神經(jīng)網(wǎng)絡(luò)模型之后，用于對(duì)神經(jīng)網(wǎng)絡(luò)模型的輸出進(jìn)行處理，得到tensor類型的數(shù)據(jù)。

有倆個(gè)參數(shù)，start_dim和end_dim，分別表示開始的維度和終止的維度，默認(rèn)值分別是1和-1，其中1表示第一維度，-1表示最后的維度。結(jié)合起來看意思就是從第一維度到最后一個(gè)維度全部給展平為張量。（注意：數(shù)據(jù)的維度是從0開始的，也就是存在第0維度，第一維度并不是真正意義上的第一個(gè)）

同理，如果我這么寫：

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

那么意思就是從第二維度開始，到第三維度全部給展平，也就是將2、3兩個(gè)維度展平。

官網(wǎng)給出的示例：

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

#開頭的代碼是注釋

整段代碼的意思是：給定一個(gè)維度為（32，1，5，5）的隨機(jī)數(shù)據(jù)。

1.先使用一次nn.Flatten()，使用默認(rèn)參數(shù)：

m = nn.Flatten()

也就是說從第一維度展平到最后一個(gè)維度，數(shù)據(jù)的維度是從0開始的，第一維度實(shí)際上是數(shù)據(jù)的第二個(gè)位置代表的維度，也就是樣例中的1。

因此進(jìn)行展平后的結(jié)果也就是[32，1×5×5]?[32，25]

2.接著再使用一次指定參數(shù)的nn.Flatten()，即

m = nn.Flatten(0, 2)

也就是說從第0維度展平到第2維度，0~2，對(duì)應(yīng)的也就是前三個(gè)維度。

因此結(jié)果就是[32×1×5，5]?[160，5]

因此進(jìn)行展平后的結(jié)果也就是[32，1*5*5]?[32，25]

示例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])

總結(jié)

到此這篇關(guān)于pytorch中nn.Flatten()函數(shù)詳解的文章就介紹到這了,更多相關(guān)pytorch nn.Flatten()函數(shù)詳解內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: