torch.nn.Linear的维度变换过程详解（有图有公式有代码）

当初在学习nn.Linear时了解到的博客都是关于一维变换的，比如输入3通道，输出6通道；又比如得到（3，4，4）的特征图，需要进行拉平为（48，）的向量，然后通过nn.Linear（48，10）得到10个输出（分类任务很常见）。

nn.Linear除了可以进行分类，主要的作用就是改变维度便于下一个卷积层或线形层的输入。

但是在实际代码中，nn.Linear的输入往往都是多维数据，一样可以正常输出。所以经过查阅手册和各个帖子，给出了自己的理解，作为笔记。

目录

一、nn.Linear函数用法

二、维度变换过程

三、全连接层的参数量与计算量

一、nn.Linear函数用法

nn.Linear 是 PyTorch 框架中的一个模块，用于实现线性层，也就是全连接层。线性层是神经网络中的基本构件，它执行一个基于矩阵乘法的线性变换，通常用于将输入数据转换为输出数据。

参数介绍：

• in_features：输入特征的数量。
• out_features：输出特征的数量。
• bias：一个布尔值，指示是否使用偏置项（默认为True）。

import torch
import torch.nn as nn

# 定义输入特征的尺寸
input_height, input_width = 4, 4
# 定义输入通道数
input_channels = 3
# 定义输出节点数
output_nodes = 5

# 创建一个随机的输入特征图，维度为[2，3，4，4]
input_data = torch.randn(2, input_channels, input_height, input_width)

# 创建一个全连接层，4 -> 5
linear_layer = nn.Linear(input_data.size(-1), output_nodes)

# 应用全连接层
output = linear_layer(input_data)

# 输出的尺寸将是 [2，3，4，5]
print("Output shape:", output.shape)

可以看到[2,3,4,4]维度的数据经过nn.Linear得到了[2,3,4,5]的数据，确实可以计算多维度。

二、维度变换过程

我查了pytorch的手册，如下：

首先通过公式可以看到nn.Linear是通过一个权重矩阵来实现维度的变化的。x是输入，A是权重矩阵，x与经过转置的权重矩阵A进行矩阵乘法，最后加上偏置项。
其次nn.Linear的输入是不限制维度的，可以看到括号中的*，其中 * 表示任意数量的附加维度，包括为空（即常见的数据拉平后只剩一个维度）。
权重矩阵维度为（out,in）,但是nn.Linear函数的用法是nn.Linear(in,out)。
最终输出的结果是（*,out）。

我画了个计算维度变换图，如下：

假设输入的数据维度为[32，3，4]，通过nn.Linear(4,2)得到[32，3，2]。这里取消偏置项。
由于在手册中权重矩阵的维度是（out,in），那么而经过转置之后就是（in,out)也就是图中的（4，2）。|
最终得到（1，2）形状的输出，准确的来说，是将（4，）形状变为（2，）。

图中可以看到单个权重矩阵有8个参数，好像不多，为什么其他帖子中都说全连接层的参数量很大呢？

三、全连接层的参数量与计算量

这一章用代码输出数据来论证，还是以输入的数据维度为[32，3，4]，通过nn.Linear(4,2)得到[32，3，2]为例，在给出代码之前先猜一下两个问题。
1.  这个过程的参数量是多少？
2. 这个过程的计算量是多少？

我在很多帖子上看到说全连接层参数量很大等等结论，于是我一开始以为参数量是32*3*4*2=768，计算量也是这么多。但是实际情况并不是（他们说的是维度拉平后再输入的情况），代码如下：

import torch
import torch.nn as nn
from thop import profile
from thop import clever_format
class MyModel(nn.Module):
    def __init__(self, input_k, output_nodes):
        super(MyModel, self).__init__()
        # 全连接层
        self.linear = nn.Linear(input_k, output_nodes, bias=False)

    def forward(self, x):
        # 应用全连接层
        x = self.linear(x)
        return x

# 定义输入特征的尺寸
input_k = 4
# 定义输入通道数
input_channels = 3
# 定义输出节点数
output_nodes = 2

# 创建一个随机的输入特征图，维度为[32，3，4]
input_data = torch.randn(32, input_channels, input_k)
# 创建一个全连接层，4 -> 2
model = MyModel(input_data.size(-1), output_nodes)
# 应用全连接层
output = model(input_data)
# 输出的尺寸将是 [32，3，2]
print("Output shape:", output.shape)

# 定义一个函数来计算模型的参数量
def count_parameters(model):
    return sum(p.numel() for p in model.parameters())

# 计算并打印模型的参数量
total_params = count_parameters(model)
print("Total parameters:", total_params)

# 使用 thop 计算 FLOPs
flops, params = profile(model.to('cuda'), (input_data.to('cuda'), ), verbose=False)
# flops 已经是浮点数
print('Total GFLOPS: %s' % flops, 'Total params: %s' % params)

关于代码我解释一下，定义了一个只包含一个线性层的模型，便于计算参数量和计算量。

可以看到输出的形状只改变了最后一个维度，从[32，3，4]变为[32，3，2]。
但是参数量却等于8，只等于一个权重矩阵，难道是最后一个维度共享权重矩阵么？
于是我利用thop库得到计算量（也就是前向过程计算了多少次），发现是768，正好等于32*3*4*2。

现在来分析一下，目前来看nn.Linear只会改变数据最后一个维度的大小。那么就不会对每个样本的所有维度都分配单独的权重，这就是手册中权重矩阵的维度是（out,in）的原因，原来一切早已注定，只是官方没解释太详细。
所以现在是数据的最后维度都是共享权重，权重参数量为4*2=8，所以参数量总和是8。
每更新一次权重参数就算8次计算（每个权重矩阵有8个参数），也就是说遍历完输入数据的维度需要32*3次，那么32*3*（4*2）=768。

计算量为768也验证了共享权重的猜想。

现在回头来看手册中的内容，就理解其中的内容了。

不足之处请大佬指出！