torch.flatten()函数

Posted 2023-04-02 one-over

**

1.第一个例子

**
torch.flatten(x)等于torch.flatten(x,0)，默认将张量拉成一维的向量，也就是说从第一维开始平坦化，也就是开始整合。
torch.flatten(x,1)代表从第二维开始平坦化。
torch.flatten(x,0,1)代表在第一维和第二维之间平坦化。

代码示例：
这里的tensor有batch，就按照有的来，直接从0开始数tensor的第几维，batch就是第0维。

**

2.第二个例子

**
具体解释

torch.flatten(input, start_dim=0, end_dim=-1)
input: 一个 tensor，即要被“推平”的 tensor。
start_dim: “推平”的起始维度。
end_dim: “推平”的结束维度。

如果我们要自己设定起始维度和结束维度呢？
我们要先来看一下 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])

我们可以看到，这里的tensor的形状是[channel,h,w]，与第一个例子不同， 这里的tensor没有batch，就按照没有的来，直接从0开始数tensor的第几维，channel就是第0维。 tensor实体最外面小括号仅作为tensor的标志，然后最外层方括号包裹，要想看tensor是几通道的(channel)，看最外层方括号里面有几个“大元素”。
具体如下：最外层的方括号内含两个元素，因此 shape 的第一个值(channels)是 2；类似地，第二层方括号里面含三个元素，shape 的第二个值(H)就是 3；最内层方括号里含四个元素，shape 的第3个值(W)就是 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])

pytorch——torch.flatten() 和 torch.nn.Flatten()

flatten()函数的作用是将tensor铺平成一维

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

• input (Tensor) – the input tensor.
• start_dim (int) – the first dim to flatten
• end_dim (int) – the last dim to flatten

start_dim和end_dim构成了整个你要选择铺平的维度范围

下面举例说明

x = torch.tensor([[1,2], [3,4], [5,6]])
x = x.flatten(0)
x
------------------------
tensor([1, 2, 3, 4, 5, 6])

对于图片数据，我们往往期望进入fc层的维度为(channels, N)这样

x = torch.tensor([[[1,2],[3,4]], [[5,6],[7,8]]])
x = x.flatten(1)
x
-------------------------
tensor([[1, 2],
        [3, 4],
        [5, 6]])

注：torch.nn.Flatten(start_dim=1, end_dim=- 1)
start_dim 默认为 1

所以在构建网络时，下面两种是等价的

class Classifier(nn.Module):
    def __init__(self):
        super(Classifier, self).__init__()
        # The arguments for commonly used modules:
        # torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0)
        # torch.nn.MaxPool2d(kernel_size, stride=None, padding=0)

        # input image size: [3, 128, 128]
        self.cnn_layers = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),

            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),

            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=4, stride=4, padding=0),
        )
        self.fc_layers = nn.Sequential(
            nn.Linear(256 * 8 * 8, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 11)
        )

    def forward(self, x):
        # input (x): [batch_size, 3, 128, 128]
        # output: [batch_size, 11]

        # Extract features by convolutional layers.
        x = self.cnn_layers(x)

        # The extracted feature map must be flatten before going to fully-connected layers.
        x = x.flatten(1)

        # The features are transformed by fully-connected layers to obtain the final logits.
        x = self.fc_layers(x)
        return x

class Classifier(nn.Module):
    def __init__(self):
        super(Classifier, self).__init__()

        self.layers = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),

            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),

            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=4, stride=4, padding=0),

            nn.Flatten(),

            nn.Linear(256 * 8 * 8, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 11)
        )

    def forward(self, x):
       
        x = self.layers(x)

        return x