前言

最近在復現(xiàn)論文中一個塊的時候需要使用到torch.flatten()這個方法，這個方法其實很簡單，但其中有一些細節(jié)可能需要注意，且有個關鍵點很容易忘記，故在此記錄以備查閱。

方法解析

flatten的中文含義為“扁平化”，具體怎么理解呢？我們可以嘗試這么理解，假設你的數(shù)據(jù)為1維數(shù)據(jù)，那么這個數(shù)據(jù)天然就已經(jīng)扁平化了，如果是2維數(shù)據(jù)，那么扁平化就是將2維數(shù)據(jù)變?yōu)?維數(shù)據(jù)，如果是3維數(shù)據(jù)，那么就要根據(jù)你自己所選擇的“扁平化程度”來進行操作，假設需要全部扁平化，那么就直接將3維數(shù)據(jù)變?yōu)?維數(shù)據(jù)，如果只需要部分扁平化，那么有一維的數(shù)據(jù)不會進行扁平操作，具體看下面的案例分析。

可以看到，torch.flatten()方法有三個參數(shù)，分別:

• input tensor：該方法的輸入
• start_dim：開始flatten的維度
• end_dim：結束flatten的維度

案例解析

• 導包

import numpy as np import torch

• 案例1 – 全部扁平化

x = np.arange(27) x = np.reshape(x, (3,3,3)) x = torch.from_numpy(x) print('before flatten', x) x = torch.flatten(x) # 默認扁平化程度為最高 print('after flatten', x)

• 案例2 – 部分扁平化

x = np.arange(27) x = np.reshape(x, (3, 3, 3)) x = torch.from_numpy(x) print('before flatten', x) x = torch.flatten(x, start_dim=0, end_dim=1) print('after flatten', x)

• 案例3 – 部分扁平化

x = np.arange(27) x = np.reshape(x, (3, 3, 3)) x = torch.from_numpy(x) print('before flatten', x) print(x.shape) x = torch.flatten(x, start_dim=1, end_dim=2) print('after flatten', x)

與50位技術專家面對面20年技術見證，附贈技術全景圖

總結

以上是生活随笔為你收集整理的PyTorch基础（15）-- torch.flatten()方法的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內容還不錯，歡迎將生活随笔推薦給好友。