關於torch.flatten的筆記


torch.nn.Flatten()可以是Sequential模型的一層,torch中定義的是:

def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None:
    super(Flatten, self).__init__()
    self.start_dim = start_dim
    self.end_dim = end_dim

def forward(self, input: Tensor) -> Tensor:
    return input.flatten(self.start_dim, self.end_dim)

因此,它是從第一維(而不是第0維)推平到最后一維。因此,它保留了batchsize,僅把每個數據推平為一維向量。

維指的就是形狀,如[2,3,4]

torch.flatten

先看函數參數:
torch.flatten(input, start_dim=0, end_dim=-1)
input: 一個 tensor,即要被“推平”的 tensor。

start_dim: “推平”的起始維度。

end_dim: “推平”的結束維度。

首先如果按照 start_dim 和 end_dim 的默認值,那么這個函數會把 input 推平成一個 shape 為 [n] 的tensor,其中 n 即 input 中元素個數。

如果我們要自己設定起始維度和結束維度呢?

我們要先來看一下 tensor 中的 shape 是怎么樣的:

t = torch.tensor([[[1, 2, 2, 1],
                   [3, 4, 4, 3],
                   [1, 2, 3, 4]],
                  [[5, 6, 6, 5],
                   [7, 8, 8, 7],
                   [5, 6, 7, 8]]])
print(t, t.shape)

運行結果:

tensor([[[1, 2, 2, 1],
         [3, 4, 4, 3],
         [1, 2, 3, 4]],

        [[5, 6, 6, 5],
         [7, 8, 8, 7],
         [5, 6, 7, 8]]])
torch.Size([2, 3, 4])

我們可以看到,最外層的方括號內含兩個元素,因此 shape 的第一個值是 2;類似地,第二層方括號里面含三個元素,shape 的第二個值就是 3;最內層方括號里含四個元素,shape 的第二個值就是 4。

直接使用flatten(默認):

>>> torch.flatten(t),torch.flatten(t).shape
(tensor([1, 2, 2, 1, 3, 4, 4, 3, 1, 2, 3, 4, 5, 6, 6, 5, 7, 8, 8, 7, 5, 6, 7, 8]), torch.Size([24]))

示例代碼:

x = torch.flatten(t, start_dim=1)
print(x, x.shape)

y = torch.flatten(t, start_dim=0, end_dim=1)
print(y, y.shape)

運行結果:

tensor([[1, 2, 2, 1, 3, 4, 4, 3, 1, 2, 3, 4],
        [5, 6, 6, 5, 7, 8, 8, 7, 5, 6, 7, 8]]) torch.Size([2, 12])
tensor([[1, 2, 2, 1],
        [3, 4, 4, 3],
        [1, 2, 3, 4],
        [5, 6, 6, 5],
        [7, 8, 8, 7],
        [5, 6, 7, 8]]) torch.Size([6, 4])

可以看到,當 start_dim = 1 而 end_dim = −1 時,它把第 1 個維度到最后一個維度全部推平合並了。而當 start_dim = 0 而 end_dim = 1 時,它把第 0 個維度到第 1 個維度全部推平合並了。

(這里注意的一點是,維度是從第 0 維開始的)

而且,pytorch中的 torch.nn.Flatten 類和 torch.Tensor.flatten 方法其實都是基於上面的 torch.flatten 函數實現的。
轉載自:https://dilthey.cnblogs.com/


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM