Python 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 函數實現的。
到此這篇關於Python torch.flatten()函數案例詳解的文章就介紹到這瞭,更多相關Python torch.flatten()函數內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet!
推薦閱讀:
- Pytorch中torch.flatten()和torch.nn.Flatten()實例詳解
- pytorch下的unsqueeze和squeeze的用法說明
- Broadcast廣播機制在Pytorch Tensor Numpy中的使用詳解
- pytorch中nn.Flatten()函數詳解及示例
- 聊聊Pytorch torch.cat與torch.stack的區別