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][n] 的tensor，其中 nn 即 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。

示例代碼：

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 = 11 而 end_dim = −1−1 時，它把第 11 個維度到最後一個維度所有推平合併了。而當 start_dim = 00 而 end_dim = 11 時，它把第 00 個維度到第 11 個維度所有推平合併了。pytorch中的 torch.nn.Flatten 類和 torch.Tensor.flatten 方法其實都是基於上面的 torch.flatten 函數實現的。

 

承接Matlab、Python和C++的編程，機器學習、計算機視覺的理論實現及輔導，本科和碩士的都可，鹹魚交易，專業回答請走知乎，詳談請聯繫QQ號757160542，非誠勿擾。

 

本文同步分享在 博客「於小勇」（CSDN）。
若有侵權，請聯繫 support@oschina.cn 刪除。
本文參與「OSC源創計劃」，歡迎正在閱讀的你也加入，一塊兒分享。