PyTorch 线性层nn.Linear底层数学原理

一、 原理代码

#coding:utf8
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
import torch
import torch.nn as nn
import numpy as np

"""
numpy手动实现模拟一个线性层
"""

#搭建一个2层的神经网络模型
#每层都是线性层
class TorchModel(nn.Module):
    def __init__(self, input_size=3, hidden_size1=5, hidden_size2=2):
        super(TorchModel, self).__init__()
        #第一件事，定义要使用哪些网络层
        self.layer1 = nn.Linear(input_size, hidden_size1)   # 3 * 5
        self.layer2 = nn.Linear(hidden_size1, hidden_size2) # 5 * 2

    def forward(self, x):
        #第二件事：定义这些网络层是如何组合的
        x = self.layer1(x)      #shape: (batch_size, input_size) -> (batch_size, hidden_size1)
        y_pred = self.layer2(x) #shape: (batch_size, hidden_size1) -> (batch_size, hidden_size2) 
        return y_pred

#自定义模型
class DiyModel:
    def __init__(self, w1, b1, w2, b2):
        self.w1 = w1
        self.b1 = b1
        self.w2 = w2
        self.b2 = b2

    def forward(self, x):
        #x: 2 * 3, w1.T: 3 * 5   = 2*5  b1: 1 * 5
        hidden = np.dot(x, self.w1.T) + self.b1 #1*5
        y_pred = np.dot(hidden, self.w2.T) + self.b2 #1*2
        return y_pred


#随便准备一个网络输入
x = np.array([[3.1, 1.3, 1.2],
              [2.1, 1.3, 13]])
#建立torch模型
torch_model = TorchModel(3, 5, 2)

print(torch_model.state_dict())

print("-----------")
#打印模型权重，权重为随机初始化
torch_model_w1 = torch_model.state_dict()["layer1.weight"].numpy()
torch_model_b1 = torch_model.state_dict()["layer1.bias"].numpy()
torch_model_w2 = torch_model.state_dict()["layer2.weight"].numpy()
torch_model_b2 = torch_model.state_dict()["layer2.bias"].numpy()
print(torch_model_w1, "torch w1 权重")
print(torch_model_b1, "torch b1 权重")
print("-----------")
print(torch_model_w2, "torch w2 权重")
print(torch_model_b2, "torch b2 权重")
print("-----------")
#使用torch模型做预测
torch_x = torch.FloatTensor(x)
y_pred = torch_model.forward(torch_x)
print("torch模型预测结果：", y_pred)



# #把torch模型权重拿过来自己实现计算过程
diy_model = DiyModel(torch_model_w1, torch_model_b1, torch_model_w2, torch_model_b2)
# #用自己的模型来预测
y_pred_diy = diy_model.forward(x)
print("diy模型预测结果：", y_pred_diy)

二、解析计算步骤：np.dot(x, self.w1.T) + self.b1：手动执行矩阵乘法（输入 × 权重转置）并加上偏置。

假设我们有：

• 输入 x ：形状为 (batch_size, input_features) 的一个二维矩阵
• 权重 w1 ：形状为 (output_features, input_features) 的二维矩阵
• 偏置 b1 ：形状为 (output_features,) 的一维向量

np.dot(x, self.w1.T) + self.b1 的执行过程如下：

1. 转置权重 ：self.w1.T 将权重转置为 (input_features, output_features)，这样矩阵乘法才能匹配维度。
2. 矩阵乘法 ：np.dot(x, self.w1.T) 得到形状 (batch_size, output_features) 的结果。
3. 加偏置 ：将 b1（形状 (output_features,)）广播到每一行，加到结果中。

---

具体例子

设：

• x 有 2 个样本，每个样本有 3 个特征：

  x = [[1, 2, 3],
       [4, 5, 6]]

• w1 有 2 个输出神经元，每个接收 3 个输入：

  w1 = [[0.1, 0.2, 0.3],   # 第一个输出神经元的权重
        [0.4, 0.5, 0.6]]   # 第二个输出神经元的权重

• b1 有 2 个偏置：

  b1 = [0.1, -0.2]

第一步：转置权重

w1.T = [[0.1, 0.4],
        [0.2, 0.5],
        [0.3, 0.6]]

第二步：矩阵乘法 np.dot(x, w1.T)

计算每个样本与转置权重的点积：

• 样本1 [1,2,3]：

  • 对第一个输出神经元：1*0.1 + 2*0.2 + 3*0.3 = 0.1+0.4+0.9 = 1.4
  • 对第二个输出神经元：1*0.4 + 2*0.5 + 3*0.6 = 0.4+1.0+1.8 = 3.2

• 样本2 [4,5,6]：

  • 对第一个输出神经元：4*0.1 + 5*0.2 + 6*0.3 = 0.4+1.0+1.8 = 3.2
  • 对第二个输出神经元：4*0.4 + 5*0.5 + 6*0.6 = 1.6+2.5+3.6 = 7.7

所以 np.dot(x, w1.T) 结果为：

[[1.4, 3.2],
 [3.2, 7.7]]

第三步：加上偏置 b1 = [0.1, -0.2]

• 第一行：[1.4+0.1, 3.2-0.2] = [1.5, 3.0]
• 第二行：[3.2+0.1, 7.7-0.2] = [3.3, 7.5]

最终结果：

[[1.5, 3.0],
 [3.3, 7.5]]

---

为什么用转置？

因为权重通常以 (输出维度, 输入维度) 形式存储（每行对应一个输出神经元），而矩阵乘法要求"输入特征数"与"权重列数"匹配。转置后变为 (输入维度, 输出维度)，使得 x (批量, 输入特征) 能直接相乘得到 (批量, 输出维度)。

这个操作等价于：

# 手动循环版本
output = np.zeros((x.shape[0], w1.shape[0]))
for i in range(x.shape[0]):          # 遍历样本
    for j in range(w1.shape[0]):     # 遍历输出神经元
        output[i,j] = np.dot(x[i], w1[j]) + b1[j]

三、代码的详细逻辑拆解：

这段代码的核心目的是验证 PyTorch 框架底层线性层（Linear Layer）的数学原理 ，并演示如何将 PyTorch 的模型权重提取出来，用纯 NumPy 手动复现完全相同的计算过程。

简单来说，它在做一件事："黑盒"拆解与"白盒"复现 。它证明了 nn.Linear 本质上就是矩阵乘法加上偏置项。

1. 环境设置

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

• 作用：这是一个常见的"补丁"。在 macOS 或某些 Linux 环境下，当 Python 同时加载了多个包含 OpenMP 库的动态库（如 NumPy 和 PyTorch）时，可能会发生冲突导致崩溃。这行代码强制允许这种重复加载，防止程序报错。

2. 定义两个模型

代码定义了两个功能相同但实现方式不同的模型：

A. TorchModel (PyTorch 标准实现)

• 结构 ：一个包含两层全连接层（Linear Layer）的神经网络。
  • layer1: 输入 3 维 -> 输出 5 维。
  • layer2: 输入 5 维 -> 输出 2 维。
• 特点 ：使用 nn.Module 类，内部自动处理权重初始化、前向传播（forward）等逻辑。
• 数学本质 ：y=xWT+by = xW^T + by=xWT+b（注意 PyTorch 中 Linear 层的权重矩阵形状是 (out_features, in_features)，所以在计算时实际上是 x⋅WTx \cdot W^Tx⋅WT）。

B. DiyModel (NumPy 手动实现)

• 结构 ：一个普通的 Python 类，没有继承 nn.Module。
• 参数 ：直接接收权重矩阵 (w1, w2) 和偏置向量 (b1, b2)。
• 核心逻辑 (forward) ：
  • np.dot(x, self.w1.T) + self.b1：手动执行矩阵乘法（输入 ×\times× 权重转置）并加上偏置。
  • 这里显式地写出了线性层的数学公式，用于验证 PyTorch 内部是如何计算的。

3. 数据准备与模型初始化

x = np.array([[3.1, 1.3, 1.2], [2.1, 1.3, 13]]) # 2个样本，每个3维
torch_model = TorchModel(3, 5, 2)

• 创建了一个具体的输入数据 x（2行3列）。
• 实例化了 PyTorch 模型。此时，layer1 和 layer2 的权重和偏置是随机初始化 的（因为没给 weight 和 bias 参数）。

4. 提取权重 (关键步骤)

torch_model_w1 = torch_model.state_dict()["layer1.weight"].numpy()
# ... 依次提取 w1, b1, w2, b2

• 目的：从 PyTorch 的"黑盒"中把具体的数值（权重矩阵和偏置向量）拿出来。
• 转换 ：.numpy() 将 PyTorch 的 Tensor 转换为 NumPy 数组，以便在 DiyModel 中使用。
• 注意维度 ：
  • layer1.weight 的形状是 (5, 3) (out, in)。
  • 在 DiyModel 的 np.dot(x, w1.T) 中，w1.T 变成了 (3, 5)，正好匹配输入 x 的列数。这验证了 PyTorch 的存储方式。

5. 预测与对比

1. PyTorch 预测 ：
   • 将 x 转为 Tensor。
   • 调用 torch_model.forward()。
   • 得到结果 y_pred。
2. NumPy 预测 ：
   • 将刚才提取的 w1, b1, w2, b2 传入 DiyModel。
   • 调用 diy_model.forward(x)。
   • 得到结果 y_pred_diy。

6. 预期结果

代码最后会打印两行结果：

• torch模型预测结果：
• diy模型预测结果：

这两组结果应该是完全一致的（在浮点数精度范围内）。

总结：这段代码想说明什么？

1. 揭示本质 ：nn.Linear 并不是什么魔法，它的底层就是 y=xWT+by = xW^T + by=xWT+b 的矩阵运算。
2. 权重互通：PyTorch 的模型参数可以无缝导出到 NumPy 进行自定义计算，反之亦然（只要维度对齐）。
3. 调试与理解：通过手动复现，开发者可以深入理解神经网络每一层的数据流向和变换过程，这对于调试模型或实现自定义算子非常有帮助。

一句话概括：这段代码展示了如何用纯数学公式（NumPy）完美复刻 PyTorch 线性层的计算过程，证明了框架底层的数学原理。