PyTorch 结构与架构分析

本文由 简悦 SimpRead 转码后整理润色。
原文地址：CSDN 博客
原文发布时间：2021-10-10 19:12:06
原文版权：遵循 CC BY-SA 4.0 协议，转载请附上原文出处链接及版权声明。

目录

• 1. PyTorch、Torch 与 torch 的关系
• 2. PyTorch 的整体结构
• 3. torch 核心模块
• 4. torchvision 模块
• 5. 从分层角度理解 PyTorch 架构
• 6. 架构六要素分析
• 7. 学习建议

1. PyTorch、Torch 与 torch 的关系

很多初学者会疑惑：明明学习的是 PyTorch，为什么代码中导入的是 torch？

可以从三个层次理解：

名称	含义	说明
Torch	早期科学计算框架	最初主要面向数值计算和机器学习，强调张量运算与 GPU 加速。
PyTorch	Python 深度学习框架	由 Facebook AI Research 推出，面向 Python 生态，继承了 Torch 的张量计算思想。
torch	PyTorch 的核心 Python 包	实际编码时通过 import torch 使用张量、自动求导、神经网络、优化器等功能。

简言之，PyTorch 是框架名称，torch 是它在 Python 中暴露出来的核心包。

2. PyTorch 的整体结构

从使用者视角看，PyTorch 可以粗略分为两部分：

• torch：PyTorch 的主体，提供张量计算、自动求导、神经网络模块、优化器、CUDA 支持、数据加载工具等能力。
• torchvision：面向计算机视觉任务的扩展库，提供常用数据集、模型结构、图像变换和视觉相关工具。

原文结构图如下：

需要注意的是，torchvision 不是 PyTorch 的全部，而是基于 torch 构建的领域扩展库。类似地，PyTorch 生态中也有面向音频、文本等任务方向的扩展工具。

3. torch 核心模块

torch 是 PyTorch 的核心。它围绕 Tensor 建立计算表达能力，并通过 nn、autograd、optim、utils 等模块组成完整的深度学习开发流程。

3.1 torch.Tensor

Tensor 是 PyTorch 中最基础、最重要的数据结构。可以把它理解为支持 GPU 加速和自动求导的多维数组。

主要作用：

• 表示标量、向量、矩阵以及更高维数据。
• 保存数据类型、形状、设备位置等信息，例如 float32、int64、CPU、CUDA。
• 支持加减乘除、矩阵乘法、索引、切片、广播、统计运算等数学操作。
• 可与自动求导系统配合，记录计算图并参与反向传播。

在深度学习任务中，输入数据、模型参数、中间特征、损失值和梯度通常都以 Tensor 的形式存在。

3.2 torch.Storage

Storage 可以理解为 Tensor 底层数据的存储区域。Tensor 描述的是“如何解释数据”，例如形状、步长、数据类型和设备；Storage 更接近“数据实际放在哪里”。

一般使用中，我们很少直接操作 Storage，更多是通过 Tensor 的接口完成设备迁移和类型转换，例如：

tensor = tensor.to("cuda")
tensor = tensor.float()

如果模型和数据要在 GPU 上计算，二者需要放在同一个设备上，否则容易出现设备不一致的错误。

3.3 torch.nn

torch.nn 是搭建神经网络的核心模块，包含模型基类、参数管理、常用网络层、损失函数和函数式接口。

子模块 / 类	作用	常见用途
nn.Module	所有神经网络模块的基类	自定义模型时继承它，并实现 forward()。
nn.Parameter	可训练参数的封装	注册权重、偏置等需要优化的参数。
nn.Sequential	顺序容器	按顺序堆叠网络层，适合简单前馈网络。
nn.Linear	全连接层	分类器、MLP、特征映射等。
nn.Conv1d/2d/3d	卷积层	处理序列、图像、体数据等。
nn.MaxPool/AvgPool	池化层	下采样、降低特征尺寸。
nn.RNN / nn.LSTM / nn.GRU	循环神经网络层	处理序列建模任务。
nn.init	参数初始化工具	对权重进行 Xavier、Kaiming 等初始化。
nn.functional	函数式接口	调用激活函数、归一化、损失函数等无状态操作。

其中，nn.Module 是最关键的抽象。自定义模型通常写成：

import torch.nn as nn
 
 
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 2)
 
    def forward(self, x):
        return self.fc(x)

继承 nn.Module 后，PyTorch 会自动管理子模块和参数，使模型能够与优化器、保存加载、训练模式切换等机制协同工作。

3.4 torch.autograd

torch.autograd 负责自动求导。它会在前向计算过程中动态构建计算图，并在调用 backward() 时根据链式法则计算梯度。

典型流程如下：

loss = criterion(output, target)
loss.backward()

需要注意：

• 现代 PyTorch 中，Variable 已经合并进 Tensor，通常不再单独使用 Variable。
• 如果一个 Tensor 设置了 requires_grad=True，PyTorch 会跟踪它参与的运算。
• 反向传播后，梯度会保存在参数的 .grad 属性中。

3.5 torch.optim

torch.optim 提供常用优化算法，例如 SGD、Adam、AdamW、RMSprop 等。

它的职责不是“求导”，而是根据 autograd 已经计算好的梯度更新模型参数。典型训练步骤是：

optimizer.zero_grad()
output = model(inputs)
loss = criterion(output, targets)
loss.backward()
optimizer.step()

其中：

• zero_grad()：清空上一轮梯度。
• backward()：计算当前损失对参数的梯度。
• step()：根据优化算法更新参数。

3.6 torch.cuda

torch.cuda 提供与 NVIDIA GPU 相关的能力，包括设备检测、显存管理、随机种子、CUDA Tensor 计算等。

常见用法：

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
inputs = inputs.to(device)

GPU 加速的前提是模型、输入数据和目标张量位于同一个设备上。

3.7 torch.utils.data

torch.utils.data 主要用于数据集组织和批量加载，是训练流程中的重要工具。

类	作用
Dataset	定义数据如何读取，通常需要实现 __len__() 和 __getitem__()。
DataLoader	按 batch 加载数据，支持 shuffle、多进程读取、自动拼接 batch 等功能。

常见训练流程是先定义 Dataset，再交给 DataLoader 生成批量数据：

loader = DataLoader(dataset, batch_size=32, shuffle=True)

4. torchvision 模块

torchvision 是面向计算机视觉任务的 PyTorch 扩展库，常用于图像分类、检测、分割等任务。

常用组成：

• torchvision.datasets：封装常见视觉数据集。
• torchvision.models：提供常用视觉模型结构和预训练权重。
• torchvision.transforms：提供图像裁剪、缩放、归一化、数据增强等转换。
• torchvision.ops：提供视觉任务中的常用算子，例如 NMS、RoI Align 等。

如果 torch 是 PyTorch 的通用计算与训练核心，那么 torchvision 可以看作计算机视觉方向的工具箱。

5. 从分层角度理解 PyTorch 架构

原文架构图如下：

为了更清晰地理解各模块之间的关系，可以把 PyTorch 划分为四层：

flowchart TD
    A[应用层：Dataset / DataLoader / torchvision] --> B[网络搭建层：nn.Module / Layers / functional]
    B --> C[优化层：autograd / optim]
    C --> D[数据与计算层：Tensor / Storage / CUDA]

5.1 数据与计算层

这一层由 Tensor、Storage、CUDA 设备支持和底层运算组成，解决“数据如何表示、存储和计算”的问题。

典型职责：

• 定义张量结构。
• 管理数据类型、形状、设备。
• 执行张量数学运算。
• 支持 CPU/GPU 计算。

5.2 网络搭建层

这一层以 torch.nn 为核心，解决“如何组织神经网络结构”的问题。

典型职责：

• 用 nn.Module 定义模型。
• 用卷积层、全连接层、归一化层、激活函数等构建网络。
• 用 Parameter 管理可训练参数。
• 用 init 初始化参数。

5.3 优化层

这一层由 autograd 和 optim 共同完成，解决“如何让模型参数变好”的问题。

典型职责：

• autograd 负责根据损失函数计算梯度。
• optim 负责根据梯度更新参数。
• 两者共同支撑反向传播和迭代训练。

5.4 应用层

这一层面向具体任务，解决“如何把数据送进模型并完成训练或推理”的问题。

典型职责：

• 用 Dataset 管理数据样本。
• 用 DataLoader 批量加载数据。
• 用 torchvision 等生态库处理图像、模型和数据增强。
• 编写训练、验证、测试和推理流程。

6. 架构六要素分析

从软件架构角度，可以用可用性、可修改性、性能、安全性、可测试性和易用性六个维度分析 PyTorch。

6.1 可用性

PyTorch 的可用性较好。常见错误通常能给出明确提示，例如张量维度不匹配、数据类型不一致、CPU/GPU 设备不一致等。对使用者来说，这类错误信息有助于快速定位问题。

6.2 可修改性

PyTorch 的模块边界较清晰。普通使用者可以通过继承 nn.Module 自定义网络，通过重写 Dataset 适配自己的数据，通过组合不同损失函数和优化器扩展训练逻辑。

如果进一步研究源码，也可以沿着 Tensor、autograd、nn、optim 等模块逐层理解，不需要一开始就进入所有底层实现。

6.3 性能

PyTorch 在 GPU 训练、张量计算和自动求导方面具备较高性能。对大多数深度学习研究和工程任务来说，其性能足以支撑模型训练与推理。

不过，最终性能仍然取决于多种因素，例如模型结构、batch size、数据加载速度、设备类型、混合精度、算子融合以及分布式训练策略。

6.4 安全性

深度学习框架的安全性不只涉及网络攻击，也包括模型文件加载、第三方数据集、依赖包、反序列化对象等风险。

在使用 PyTorch 时，应尽量从可信来源加载模型和数据，谨慎执行未知来源的脚本或模型文件，并在工程环境中做好依赖版本管理。

6.5 可测试性

PyTorch 的模块化设计有利于测试。模型、数据集、损失函数、训练循环都可以拆开验证。

常见测试方式包括：

• 检查模型前向传播输出形状。
• 用小批量数据验证训练循环是否能跑通。
• 检查 loss 是否下降。
• 为自定义 Dataset、数据增强、损失函数编写单元测试。

6.6 易用性

PyTorch 的易用性是它广泛流行的重要原因之一。它采用动态图机制，代码风格接近普通 Python，调试体验自然，适合研究原型开发，也能扩展到工程训练和部署场景。

对于有 Python 基础和深度学习基础的使用者来说，通常可以较快上手。

7. 学习建议

理解 PyTorch 架构时，可以按下面顺序学习：

1. 先掌握 Tensor：形状、数据类型、设备、广播、索引、矩阵运算。
2. 再学习 autograd：理解 requires_grad、计算图、backward() 和 .grad。
3. 然后学习 nn.Module：掌握自定义模型、参数注册、前向传播。
4. 接着学习 optim：理解梯度清零、反向传播、参数更新的训练闭环。
5. 最后学习 Dataset 和 DataLoader：把真实数据稳定、高效地送入模型。

完整训练流程可以概括为：

准备数据 -> 构建模型 -> 定义损失函数 -> 选择优化器 -> 前向传播 -> 计算损失 -> 反向传播 -> 更新参数

把这条主线串起来后，再深入研究 torchvision、混合精度、分布式训练、模型保存加载和部署，会更容易建立系统性的理解。