pytorch中tensor的形状 [n,1] 与 [n] 的区别

本文已参与「新人创作礼」活动，一起开启掘金创作之路。

首先看一下 [n,1] 与 [n] 的区别，最根本的区别是形状为 [n,1] 的tensor是二维的，而形状为 [n] 的tensor是一维的。

一个简答的例子就能直接明白二者的区别：

import torch

a = torch.randn(3, 1)
print(f"a.shape = {a.shape}")     # a.shape = torch.Size([3, 1])
print(f"a.ndim = {a.ndim}")       # a.ndim = 2
print(f"a = {a}")                 # a = tensor([[ 0.0685], [-1.3986], [-0.5098]])

b = a.view(-1)
print(f"b.shape = {b.shape}")     # b.shape = torch.Size([3])
print(f"b.ndim = {b.ndim}")       # b.ndim = 1
print(f"b = {b}")                 # b = tensor([ 0.0685, -1.3986, -0.5098])

那么什么时候会用到形状为 [n,1] 的 tensor呢？什么时候需要将其转化为形状为 [n] 的 tensor呢？

一个简单的例子是：在一个二分类判别器网络中，网络输入一张图片，输出0-1之间的一个数（概率）。那么输入 batch_size 张图片，网络输出 batch_size 个数，则 output.shape= [batch_size,1] ，代码如下：

import torch
import torch.nn as nn


class Discriminator(nn.Module):
    def __init__(self, img_dim):
        super().__init__()
        self.disc = nn.Sequential(
            nn.Linear(in_features=img_dim, out_features=128),
            nn.LeakyReLU(0.01),
            nn.Linear(128, 1),
            nn.Sigmoid(),   # 将输出值映射到0-1之间
        )

    def forward(self, x):
        return self.disc(x)


if __name__ == '__main__':
    batch_size = 32
    img_dim = 28 * 28 * 1
    input = torch.randn(batch_size, img_dim)
    disc = Discriminator(img_dim)
    output = disc(input)
    print(f"output.shape = {output.shape}")     # output.shape = torch.Size([32, 1])

如果真实的label均为1，那么需要计算 output 和 label 之间的 loss , 需要使用 torch.ones_like(output)) 函数生成一个形状与 output 相同的 元素均为1的tensor

计算 loss：

import torch
import torch.nn as nn


class Discriminator(nn.Module):
    def __init__(self, img_dim):
        super().__init__()
        self.disc = nn.Sequential(
            nn.Linear(in_features=img_dim, out_features=128),
            nn.LeakyReLU(0.01),
            nn.Linear(128, 1),
            nn.Sigmoid(),   # 将输出值映射到0-1之间
        )

    def forward(self, x):
        return self.disc(x)


if __name__ == '__main__':
    batch_size = 32
    img_dim = 28 * 28 * 1
    input = torch.randn(batch_size, img_dim)

    disc = Discriminator(img_dim)
    output = disc(input)
    print(f"output.shape = {output.shape}")     # output.shape = torch.Size([32, 1])
    
    # 计算loss
    criterion = nn.BCELoss()
    loss1 = criterion(output, torch.ones_like(output))
    print(f"loss1 = {loss1}")        # loss1 = 0.714613676071167
    loss2 = criterion(output.view(-1), torch.ones_like(output.view(-1)))
    print(f"loss2 = {loss2}")        # loss2 = 0.714613676071167

可以看到，使用 [32, 1] 和 [32]的 tensor 计算 loss 的结果是完全相同的。