torch.flatten vs torch.nn.Flatten 的區別

torch.flatten 和 torch.nn.Flatten 都用於把多維Tensor展平(flatten), 區別是：

• torch.flatten是函數，使用前無需先實例化，默認從第0維開始展平，通用化好

torch.flatten(input, start_dim=0, end_dim=- 1)

• torch.nn.Flatten是類，使用前需要先實例化，由於其在torch.nn模塊中，默認專門處理神經網數據的展平，而神經網絡數據通常第0維是Batch_Size, Batch_Size無需展平，所以其默認從第1維開始展平。

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

測試範例程序如下：

import torch

input_tensor = torch.randn(32, 4, 5, 5)
m = torch.nn.Flatten() #實例化Flatten
output1 = m(input_tensor)
print(output1.shape)
output2 = torch.flatten(input_tensor)
print(output2.shape)

運行結果如下：

torch.Size([32, 100])
torch.Size([3200])

另外，torch.nn.Flatten適合作爲一個“神經網絡層”，加入神經網絡中，範例：

def _create_fcs(self, split_size, num_boxes, num_classes):
        S, B, C = split_size, num_boxes, num_classes
        return nn.Sequential(
            nn.Flatten(),
            nn.Linear(1024 * S * S, 4096), 
            # Usually, dropout is placed on the fully connected layers only
            # A rule of thumb is to set the keep probability (1 - drop probability) to 0.5 when dropout is applied to fully connected layers
            # https://stackoverflow.com/questions/46841362/where-dropout-should-be-inserted-fully-connected-layer-convolutional-layer
            nn.Dropout(0.5),
            nn.LeakyReLU(0.1),
            # The predictions are encoded as an S × S × (B ∗ 5 + C) tensor
            nn.Linear(4096, S * S * (B * 5 + C)), # 7*7*(2*5+20)=1470
        )