Pytorch 源码学习笔记梳理

一、PyTorch的主要模块

1. torch 模块

torch模块包含一些激活函数、Pytorch张量的运算操作、及特定形状张量的生成。

2. torch.Tensor 模块

torch.Tensor模块定义了torch中张量的类型。同时，张量的类中也包含着一系列的方法，可以操作当前的张量，从而返回新的张量或更改当前的张量。
torch.Storage 模块则负责torch.Tensor底层的数据存储，为张量分配连续的一维内存地址。
这里需要提到一点，如果张量的某个类方法会返回张量，安装PyTorch中的命名规则，如果方法后缀带下划线，则会修改张量本身数据，类似pandas中的inplace=True。如Tensor.add和Tensor.add_。

3. torch.sparse 模块

torch.sparse模块定义了稀疏张量。构造稀疏张量采用的是COO格式（Coordinate），主要方法是用一个长整形定义非零元素的位置，用浮点数张量定义对应非零元素的值。

4. torch.cuda 模块

torch.cuda 模块定义了与CUDA运算相关的一系列函数，包括但不限于检查系统的CUDA是否可用，当前进程对应的GPU序号，清除GPU上的缓存，设置GPU的计算流（Stream），同步GPU上执行的所有核函数（Kernel）等。

5. torch.nn 模块

torch.nn是一个非常重要的模块，是PyTorch神经网络模块化的核心。定义了一系列神经网络结构，包括卷积层nn.ConvNd(N=1,2,3)和线性层（全连接层）nn.Linear等等。当构建深度学习模型时，可以通过继承nn.Module类并重写forward方法来实现一个新的神经网络。
此外，torch.nn中也定义了一系列的损失函数，如torch.nn.MSELoss, torch.nn.CrossEntropyLoss等。

6. torch.nn.functional 模块

torch.nn.functional是PyTorch的函数模块，定义了一些与神经网络相关的函数，包括卷积函数、池化函数、激活函数、dropout函数、损失函数等等。torch.nn中定义的模块一般会调用torch.nn.functional里的函数。

7. torch.nn.init 模块

torch.nn.init 模块定义了神经网络权重的初始化方法。如均匀初始化torch.nn.init.uniform_和正太分布归一化torch.nn.init.normal_等

8. torch.optim 模块

torch.optim 模块定义了一系列的优化算法，如torch.optim.SGD, torch.optim.AdaGrad, torch.optim.Adam等。
同时，这个模块还包含学习率衰减算法的子模块，即torch.optim.lr_scheduler。

9. torch.autograd 模块

torch.autograd 模块提供自动微分算法。其中的torch.autograd.backward 函数，主要用于在求得损失函数之后进行反向梯度传播，torch.autograd.grad函数用于一个标量张量（一维张量）对另一个张量求导，以及在代码中设置不参与求导的部分。另外，这个模块还内置了数值梯度功能、检查自动微分引擎是否输出正确结果等功能。

10. torch.distributed 模块

torch.distributed 模块提供分布式计算方法，其主要支持的后端有MPI，Gloo，NCCL三种。
PyTorch的分布式工作原理主要是启动多个并行的进程，每个进程都拥有一个模型的备份，然后输入不同的训练数据到多个并行的进程，计算损失函数，每个进程独立地做反向传播，最后对所有进程权重张量的梯度做归约（Reduce）。用到后端的部分主要是数据的广播（Broadcast）和数据的收集（Gather），其中，前者是把数据从一个节点（进程）传播到另一个节点（进程），比如归约后梯度张量的传播，后者则是把数据从其他节点转移到当前节点，比如把梯度张量从其他节点转移到某个特定的节点，然后对所有的张量求平均。
torch.distributed 模块还提供了一些启动方式来启动多个进程，包括但不限于通过网络（TCP）、通过环境变量、通过共享文件等。

11. torch.distributions 模块

torch.distributions 模块提供了一系列的概率分布函数和采样函数。

12. torch.hub 模块

torch.hub 模块提供了一系列预训练的模型供用户使用。

13. torch.jit 模块

torch.jit是即时编译器（Just-In-Time Compiler, JIT）模块。这个模块的意义是把动态图转换成可以优化和序列化的静态图，其主要工作原理是通过预先定义好的张量，追踪整个动态图的构建过程，得到最终构建出来的动态图，然后转换为静态图。通过JIT得到的静态图可以保存，并且被其他前端支持。另外，JIT也可以用来生成其他格式的神经网络描述文件，如ONNX。
需要注意的是，torch.jit支持两种模式，即脚本模式（ScriptModule）和追踪模式（Tracing）。他们都能构建静态图，区别在于前者支持控制流，后者不支持，但是前者支持的神经网络模块比后者少，比如脚本模式不支持torch.nn.GRU。

14. torch.linalg 模块

torch.linalg 模块包含一系列常见的线性代数运算

15. torch.special 模块

torch.special 模块包含一系列的针对张量的科学运算、特殊函数等

16. torch.package 模块

torch.package 模块用于封装PyTorch代码和模型，以便在不同机器上运行

17. torch.profiler 模块

torch.profiler 模块用来在训练和预测期间收集性能指标，研究设备内核活动并可视化执行跟踪，从而帮助优化模型训练过程，优化模型性能。

18. torch.onnx 模块

torch.onnx 模块用于导出和载入ONNX格式的深度学习模型描述文件。ONNX格式的存在是为了方便不同深度学习框架之间交换模型。

19. torch.random 模块

torch.random 模块提供了一系列方法来保存和设置随机数生成器的状态，包括使用get_rng_state 函数获取当前随机数生成器状态，set_rng_state 函数设置当前随机数生成器状态，并且可以用manual_seed 函数来设置随机种子，也可以使用initial_seed 函数来得到程序初始的随机种子。
因为神经网络的训练是一个随机过程，包括数据的输入、权重的初始化都具有一定的随机性。设置一个统一的随机种子可以有效地帮助我们测试不同结构神经网络的表现，有助于调试神经网络的结构。

20. torch.utils.bottleneck 模块

torch.utils.bottleneck 模块用来检查深度学习模型中各个组件的运行时间，从而找到导致性能瓶颈的地方，通过优化瓶颈处的运行时间来优化整个深度学习模型的性能。

21. torch.utils.checkpoint 模块

torch.utils.checkpoint 模块可以记录反向传播过程中中间数据的计算过程，然后丢弃这些中间数据，等需要用到的时候再重新计算。这个模块设计的核心思想是以计算时间换内存空间，从而节约深度学习使用的内存。

22. torch.utils.cpp_extension 模块

torch.utils.cpp_extension 模块定义了PyTorch的C++扩展，其主要包含两个类：

• CppExtension 定义了使用C++来编写的扩展模块的源代码相关信息
• CUDAExtension 则定义了C++/CUDA 编写的扩展模块的源代码相关信息

在底层上，这个扩展模块使用了pybind11，保持了接口的轻量性并使得PyTorch 易于被扩展。

23. torch.utils.data 模块

torch.utils.data 模块引入了数据集（Dataset）和数据载入器（DataLoader）的概念。数据载入器通过对数据集的包装，可以对数据集进行随机排列（Shuffle）和采样（Sample），得到一系列打乱数据顺序的迷你批次。

24. torch.utils.tensorboard 模块

torch.utils.tensorboard 模块是 PyTorch 对 TensorBoard 数据可视化工具的支持，能够显示深度学习模型在训练过程中损失函数、张量权重的直方图，以及模型训练过程中输出的文本、图像、视频等。
TensorBoard 的功能非常强大，而且是基于可交互的动态网页设计的，使用者可以通过预先提供的一系列功能来输出特定的训练过程的细节，有助于方便地调试深度学习模型。