PyTorch的底层是基于C++实现的，核心部分包括计算图（Computation Graph）、张量计算（Tensor Computation）、自动微分（Autograd）、运行时调度（Runtime Execution）、内存管理（Memory Management）以及CUDA优化等。PyTorch主要利用了C++（特别是ATen、c10等库）来高效地进行计算，并通过Python绑定提供用户友好的接口。
下面，我从核心模块入手，深入讲解PyTorch的底层原理，并涉及相关的C++代码分析。
1. PyTorch的核心组件
PyTorch底层主要由以下几个部分组成：
- ATen（Tensor 库） - 负责张量的基本操作，底层封装了Eigen、MKL、CUDA等计算库。
- Autograd（自动微分引擎） - 计算梯度，构建计算图。
- c10（核心组件） - 主要提供调度和设备管理（Device Management）。
- JIT（TorchScript 运行时） - 用于序列化和优化计算图。
- Dispatcher（调度器） - 负责动态调用不同的计算后端（CPU、CUDA、XLA）。
- Memory Management（内存管理） - 主要涉及PyTorch的内存池（Caching Allocator）。
2. PyTorch Tensor 底层（ATen库）
PyTorch 的 Tensor 主要由 ATen（A Tensor Library） 实现，它是一个轻量级的张量库，提供跨 CPU/GPU 统一的接口。ATen 采用 C++ 实现，并对外暴露了 C++ API。
2.1 ATen 的基本架构
- 核心数据结构：
TensorImpl - 操作符封装：
Tensor类 - 后端管理：
Dispatch - 计算后端：CPU（Eigen, MKL） / CUDA（cuBLAS, cuDNN）
ATen 的 Tensor 底层数据结构为 TensorImpl，它存储张量的核心信息：
cpp
struct TensorImpl {
caffe2::TypeMeta dtype_; // 数据类型
at::Storage storage_; // 存储数据
at::Device device_; // 设备（CPU / CUDA）
int64_t sizes_[N]; // 维度信息
int64_t strides_[N]; // 步长信息
};关键点：
storage_是 PyTorch 共享的底层存储，多个 Tensor 可以共享同一块数据（通过view()等方式）。sizes_和strides_负责存储张量形状和内存步长（对于transpose()操作尤为重要）。dtype_定义数据类型（float,int,double等）。
2.2 创建一个 PyTorch Tensor
当你执行：
python
import torch
x = torch.randn(3, 3)底层的 C++ 代码大致如下：
cpp
Tensor at::empty(IntArrayRef size, const TensorOptions& options) {
auto device = options.device();
auto dtype = options.dtype();
return at::detail::empty_cuda(size, dtype) // GPU
: at::detail::empty_cpu(size, dtype); // CPU
}具体调用了 empty_cpu 或 empty_cuda，如果是 CPU：
cpp
Tensor empty_cpu(IntArrayRef size, const TensorOptions& options) {
auto storage = at::Storage(size.numel() * sizeof(float));
return at::Tensor(c10::make_intrusive<TensorImpl>(storage, at::DeviceType::CPU));
}这个 Storage 是 PyTorch 的底层内存管理结构，负责存储数据。
3. Autograd（自动微分）
PyTorch 采用 动态计算图，即每次计算都会构建一张新的计算图，而不像 TensorFlow 那样静态编译计算图。
3.1 计算图（Computation Graph）
PyTorch 的 Autograd 采用了一种基于 反向模式自动微分 （Reverse-mode Autograd）的机制。计算图由 Node 组成，每个 Tensor 记录了其生成操作 grad_fn。
python
import torch
x = torch.randn(3, requires_grad=True)
y = x * 2
z = y.mean()
z.backward()在底层，PyTorch 构建了如下的计算图：
x
\
*2 (MulBackward)
\
mean (MeanBackward)
|
grad每个计算操作都会产生一个 Function 节点：
cpp
struct Node {
virtual void apply() = 0; // 反向传播的实现
std::vector<Tensor> next_edges_; // 指向下一层的边
};关键 C++ 代码：
cpp
struct MulBackward : public Node {
void apply() override {
grad_input = grad_output * 2;
}
};计算 z.backward() 时，PyTorch 会从 z 的 grad_fn 开始，依次调用 apply() 计算梯度，并将梯度回传。
4. Dispatcher（调度器）
PyTorch 使用 动态调度 机制来选择合适的后端（CPU / CUDA / XLA）。调度器的作用是将高层 API 请求映射到不同设备的实现上。
4.1 调度器结构
在 PyTorch 的 C++ 代码中，调度器的核心组件是 Dispatcher：
cpp
class Dispatcher {
public:
void registerKernel(const std::string& op_name, Kernel kernel);
Kernel lookupKernel(const std::string& op_name, DeviceType device);
};当你执行 torch.add(x, y) 时，PyTorch 需要判断 x 是否在 GPU 上，然后调度到不同的实现：
cpp
Tensor add(const Tensor& a, const Tensor& b) {
if (a.device().is_cuda()) {
return add_cuda(a, b);
} else {
return add_cpu(a, b);
}
}这就是 PyTorch 高效支持多后端计算的核心机制。
5. CUDA 加速
PyTorch 使用 cuBLAS, cuDNN 等库进行 CUDA 加速，并使用 CUDA Graphs 提高计算效率。
5.1 CUDA 内存管理
PyTorch 采用了一种 缓存分配器（Caching Allocator），防止 GPU 内存碎片化：
cpp
void* THCCachingAllocator::malloc(size_t size) {
void* ptr = cudaMalloc(size);
cache.insert(ptr);
return ptr;
}这样可以避免频繁 cudaMalloc() 带来的性能开销。
总结
- PyTorch 的 Tensor 由 ATen 实现，TensorImpl 是核心数据结构。
- Autograd 采用动态计算图，每个计算节点记录
grad_fn。 - Dispatcher 负责自动调度不同计算后端（CPU / CUDA）。
- CUDA 采用缓存分配器，提高 GPU 内存利用率。
PyTorch 的底层代码非常复杂，但掌握这些核心概念后，你可以更深入地理解 PyTorch 的实现原理！
PyTorch 的自动微分（Autograd）系统是基于 计算图（Computation Graph） 和 反向模式自动微分（Reverse-mode Autodiff） 来实现的。核心思想是 动态计算图 ：在前向计算时，PyTorch 记录运算操作，并在反向传播时利用链式法则计算梯度。接下来，我们深入讲解 如何实现自动微分 ，并剖析 PyTorch C++ 源码。
1. 计算图（Computation Graph）
计算图是一种 有向无环图（DAG），每个节点代表一个计算操作，每条边代表数据的流动（即 Tensor）。
示例代码
python
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 # y = x^2
z = y + 3 # z = x^2 + 3
z.backward() # 计算 dz/dx
print(x.grad) # 输出 4计算图结构
x
\
^2 (PowBackward)
\
+3 (AddBackward)
|
dz/dx每次计算都会动态创建 Function 节点（如 PowBackward 和 AddBackward），并存储在 grad_fn 中。
如何自动构建计算图？
- 每个 Tensor 记录它的计算来源（grad_fn）
- 每个操作都会注册反向传播函数
- 计算 z.backward() 时，沿着计算图反向传播
2. 计算图的底层实现
2.1 关键数据结构
在 PyTorch 源码中，计算图由 Node（计算节点）和 Edge（计算依赖）组成。
计算节点 (Node)
cpp
struct Node {
virtual void apply() = 0; // 反向传播计算
std::vector<Tensor> next_edges_; // 依赖的梯度计算
};Node 是所有计算操作的基类，每个操作（如 MulBackward）都会继承它并实现 apply() 方法。
计算边 (Edge)
cpp
struct Edge {
Node* function; // 指向计算节点
int input_nr; // 输入编号（多个输入时使用）
};每个 Tensor 都会有一个 grad_fn 指向其生成的 Node，从而形成计算图。
2.2 计算图的构建
在 Tensor 计算时，PyTorch 自动构建计算图。例如：
cpp
Tensor pow(const Tensor& self, double exponent) {
auto result = at::empty(self.sizes(), self.options());
at::native::pow_out(result, self, exponent);
if (self.requires_grad()) {
auto grad_fn = std::make_shared<PowBackward>(self, exponent);
result.set_grad_fn(grad_fn);
}
return result;
}- 创建
PowBackward计算节点 - 如果
self需要梯度，关联grad_fn - 最终在
z.backward()时，反向遍历grad_fn计算梯度
3. 反向传播（Backpropagation）
当 z.backward() 运行时，PyTorch 从 grad_fn 开始，沿计算图反向传播。
3.1 反向传播的链式法则
根据链式法则：
\\frac{dz}{dx} = \\frac{dz}{dy} \\times \\frac{dy}{dx}
PyTorch 遍历计算图，依次计算梯度。
3.2 关键代码
PyTorch 反向传播的核心实现位于 torch/csrc/autograd/engine.cpp：
梯度计算调度
cpp
void Engine::execute(Node* graph_root) {
std::deque<Node*> queue;
queue.push_back(graph_root);
while (!queue.empty()) {
Node* node = queue.front();
queue.pop_front();
// 调用 apply 计算梯度
node->apply();
for (Edge& edge : node->next_edges_) {
queue.push_back(edge.function);
}
}
}- 使用队列进行反向遍历
- 每个
Node调用apply()计算梯度 - 遍历
next_edges_继续传播
4. Autograd 具体实现
4.1 计算 y = x^2 的梯度
cpp
struct PowBackward : public Node {
Tensor input_;
double exponent_;
PowBackward(const Tensor& input, double exponent)
: input_(input), exponent_(exponent) {}
void apply() override {
// 计算 dy/dx = 2 * x
Tensor grad_input = exponent_ * input_.pow(exponent_ - 1);
input_.grad().add_(grad_input);
}
};当 z.backward() 被调用时：
PowBackward计算dy/dx = 2xAddBackward计算dz/dy = 1- 乘法链式法则计算
dz/dx = dz/dy * dy/dx
4.2 计算 z = y + 3 的梯度
cpp
struct AddBackward : public Node {
Tensor input_;
AddBackward(const Tensor& input) : input_(input) {}
void apply() override {
input_.grad().add_(1); // dz/dy = 1
}
};最终：
\\frac{dz}{dx} = \\frac{dz}{dy} \\times \\frac{dy}{dx} = 1 \\times 2x = 2x
如果 x = 2，则 x.grad = 4。
5. Autograd 计算优化
5.1 梯度缓存
为了避免重复计算，PyTorch 缓存已计算的梯度：
cpp
void AccumulateGrad::apply() {
if (!variable_.grad().defined()) {
variable_.mutable_grad() = grad_input;
} else {
variable_.mutable_grad() += grad_input;
}
}如果 x 参与多个计算，它的梯度会累积，而不会覆盖。
5.2 动态计算图
PyTorch 采用 即时构建计算图 （而不是 TensorFlow 的静态图），这样每次 forward 计算都能自动记录计算路径。
5.3 多线程并行计算
PyTorch 在 Autograd 计算时采用 多线程调度：
cpp
std::thread worker([&] { node->apply(); });
worker.join();这提升了梯度计算的效率。
6. 总结
PyTorch 的自动微分系统基于 计算图 + 反向传播，其实现方式如下：
- 每个
Tensor记录grad_fn，构建动态计算图 - 计算梯度时，调用
grad_fn依次执行反向传播 - 每个
Node（如MulBackward、AddBackward）实现apply()方法 - 梯度累积、缓存优化，提升计算性能
- 多线程调度提升 Autograd 性能
这种机制使得 PyTorch 的自动微分既 灵活 又 高效 ，支持 动态图计算，而不像 TensorFlow 需要提前定义计算图。掌握这些底层实现后，你就能更深入地理解 PyTorch 的计算原理，并优化深度学习训练！
在 PyTorch 的 Autograd Engine 中，队列（queue）的使用方式决定了梯度计算的顺序 ，但 PyTorch 主要使用的是 双端队列（deque） ，可以既支持 FIFO（先进先出） 也可以支持 LIFO（后进先出）。
1. 先进先出（FIFO） vs. 后进先出（LIFO）
先进先出（FIFO, First In First Out）:
- 典型数据结构：队列（queue）
- 处理方式：先入计算图的节点，先进行反向传播
- 适用于 BFS（广度优先遍历）
后进先出（LIFO, Last In First Out）:
- 典型数据结构：栈（stack）
- 处理方式：后入计算图的节点，先进行反向传播
- 适用于 DFS（深度优先遍历）
PyTorch 在 Autograd 反向传播中使用 双端队列（std::deque） ，但它的核心执行模式实际上是 深度优先遍历（DFS）+ 拓扑排序 ，即 LIFO（后进先出）。
2. PyTorch 反向传播的执行方式
在 PyTorch 的 Autograd 计算过程中：
- 计算图是 有向无环图（DAG） ，反向传播时按照 拓扑排序（topological ordering） 遍历。
- PyTorch 默认使用后进先出（LIFO）策略 ，即 先计算最后一个操作的梯度，然后再回溯计算前面的梯度。
- 梯度计算按依赖关系进行 ，即 当某个节点所有子节点计算完成后，才能计算该节点的梯度。
3. PyTorch 反向传播队列的实现
在 torch/csrc/autograd/engine.cpp 中，PyTorch 使用 std::deque<Node*> 作为计算节点的队列：
cpp
void Engine::execute(Node* graph_root) {
std::deque<Node*> queue;
queue.push_back(graph_root); // 把根节点放入队列
while (!queue.empty()) {
Node* node = queue.back(); // 取出队列最后一个节点 (LIFO)
queue.pop_back(); // 删除该节点
// 计算该节点的梯度
node->apply();
// 遍历所有子节点，加入队列
for (Edge& edge : node->next_edges_) {
queue.push_back(edge.function);
}
}
}核心点
queue.back()取出最后一个节点，意味着 LIFO（后进先出） ，即 深度优先遍历（DFS）。pop_back()删除最后一个节点，保证先计算最近的操作梯度。push_back()把子节点放入队列，确保按 拓扑顺序反向传播。
4. 为什么 PyTorch 选择 LIFO（后进先出）？
-
符合链式法则
- 反向传播计算梯度时，必须先计算输出的梯度，才能计算输入的梯度：
\\frac{dz}{dx} = \\frac{dz}{dy} \\cdot \\frac{dy}{dx}
z.backward()先计算dz/dy，再计算dy/dx，所以采用 LIFO 方式，保证 先计算最近的操作，再计算前面的操作。
- 反向传播计算梯度时，必须先计算输出的梯度，才能计算输入的梯度：
-
高效递归计算
- LIFO 天然符合递归计算方式 ，与 栈结构（stack） 类似，计算 后入先出的梯度计算。
-
拓扑排序的自然选择
- 反向传播要求 从叶子节点到根节点 ，LIFO 方式符合 拓扑排序的逆序计算，可以正确地完成反向计算。
5. PyTorch Autograd 计算过程示例
假设：
python
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 # y = x^2
z = y + 3 # z = x^2 + 3
z.backward() # 计算 dz/dx计算图如下：
x
\
^2 (PowBackward)
\
+3 (AddBackward)
|
dz/dx反向传播步骤
z.backward()触发反向传播，从AddBackward开始。AddBackward计算dz/dy = 1，然后传递给PowBackward。PowBackward计算dy/dx = 2x，最终得到dz/dx = 2x * 1 = 4。- 计算完成。
LIFO 顺序
1. 取出 AddBackward (z = y + 3)
2. 计算 dz/dy = 1
3. 取出 PowBackward (y = x^2)
4. 计算 dy/dx = 2x
5. 计算 dz/dx = dz/dy * dy/dx这符合 后进先出（LIFO）+ 深度优先搜索（DFS） 计算方式。
6. 总结
| 策略 | 数据结构 | 遍历方式 | 适用场景 |
|---|---|---|---|
| FIFO（先进先出） | 队列（queue） | 广度优先遍历（BFS） | 不适用于梯度计算 |
| LIFO（后进先出） | 栈（stack）/双端队列（deque） | 深度优先遍历（DFS） | 适用于 Autograd 反向传播 |
PyTorch Autograd 采用 LIFO 方式的原因
- 符合链式法则的计算顺序
- 拓扑排序反向遍历，保证计算正确
- 深度优先搜索（DFS），提高计算效率
- 递归计算适配栈结构，便于优化
PyTorch 使用双端队列（deque），但以 LIFO 方式运行，确保反向传播能正确计算梯度。希望这个深入解析能帮你更好地理解 PyTorch 自动微分的底层实现！ 🚀
一般来说，在深度学习项目中，大多数工程师和研究人员不会直接修改 PyTorch 的底层 C++ 代码，除非有以下几种情况：
1. 你在开发一个新的深度学习库或框架
如果你正在做类似 PyTorch、TensorFlow、JAX 这样的大规模机器学习框架的底层开发，比如优化计算图执行、设计新的自动微分系统、改进 GPU 计算性能等，你可能需要深入修改 PyTorch 的 C++ 代码。例如：
- 开发新的后端计算引擎（如 XLA、MLIR）。
- 优化 Autograd 计算图，减少内存占用或加速梯度计算。
- 改进调度系统，让 PyTorch 在 CPU/GPU/TPU 上更高效地执行。
🛠 现实案例
- Facebook (Meta) 的 PyTorch 团队不断优化
ATen以提升 CPU/GPU 计算性能。 - Google 在 JAX 中设计了基于
XLA的 Autograd 引擎，使得梯度计算更高效。
2. 你在做 PyTorch 内核优化（高性能计算）
如果你在一个涉及 大规模分布式训练、低级优化、编写 CUDA 内核 的项目中，比如：
- 你需要 实现新的 CUDA 加速算子，让某个算子比 PyTorch 官方实现更快。
- 你在做 极端优化（low-level tuning），比如减少 PyTorch 计算时的 CPU/GPU 内存拷贝。
🔍 常见情况
-
实现新的算子（Custom Operator）
cppat::Tensor my_custom_op(const at::Tensor& input) { return input * 2; }然后在 Python 里注册：
pythonimport torch from torch.utils.cpp_extension import load my_ext = load(name="my_ext", sources=["my_op.cpp"]) x = torch.randn(10) print(my_ext.my_custom_op(x)) -
优化 PyTorch 的 Tensor 调度
- PyTorch 通过 调度器（Dispatcher） 自动选择最优的计算后端，但有时可以手动优化它：
cppTORCH_LIBRARY(my_ops, m) { m.def("my_custom_op", &my_custom_op); }
3. 你在做分布式深度学习（比如优化 PyTorch DDP）
如果你在做大规模 分布式深度学习（Distributed Deep Learning） ，你可能会修改 PyTorch 的 分布式训练底层代码（DistributedDataParallel, DDP）：
- 优化通信效率：减少 GPU 之间的梯度同步开销（如 Overlapping Computation with Communication）。
- 实现新的 AllReduce 算法：优化参数聚合。
- 深度优化 NCCL / MPI 通信：加速 GPU 之间的数据传输。
📌 案例
-
优化 PyTorch DDP 以减少 GPU 之间的通信开销：
cpp// DDP 的梯度同步逻辑 void sync_gradients() { ncclAllReduce(...); }
4. 你在做低级硬件优化（CUDA, TPU, XLA）
如果你需要让 PyTorch 运行在 新的硬件架构（比如定制 AI 加速器、TPU、专用 FPGA），你可能需要修改底层代码：
- 增加对新硬件的支持（比如让 PyTorch 支持你的 AI 芯片）。
- 优化 CUDA / HIP / ROCm 内核，实现更高效的计算。
- 修改 PyTorch 的内存管理策略，提高 GPU 内存复用率。
🔧 现实案例
- NVIDIA 在 PyTorch 里优化了
cudnn的算子执行方式，使得 PyTorch 在 A100 GPU 上运行得更快。 - Google 在 PyTorch-XLA 里增加了对 TPU 的支持，使 PyTorch 也能运行在 Google Cloud TPU 上。
5. 你在做自动微分（Autograd）系统研究
如果你在研究 新的自动微分算法（比如高阶梯度计算、低内存反向传播等），你可能会改写 PyTorch Autograd 的底层实现：
- 实现高阶梯度计算（如梯度检查点 Gradient Checkpointing）。
- 优化计算图存储方式，减少显存占用。
- 设计新的反向传播策略，比如更快的 Hessian 计算。
📌 案例
- DeepMind 在 JAX 里实现了高效的
Forward-Mode Autograd，比 PyTorch 更适合计算二阶梯度。 - Facebook AI 研究人员优化了 PyTorch 的
Autograd Engine，减少了计算图存储占用。
6. 什么时候不用改 PyTorch 底层？
如果你是 一般的深度学习算法工程师或研究员 ，你 大概率不会修改 PyTorch 底层代码 ，因为：
✅ PyTorch 提供了 Python API ，你可以通过 torch.nn、torch.optim、torch.autograd 进行训练和梯度计算，无需关心 C++ 层面。
✅ PyTorch 已有丰富的 GPU 优化 ，很多时候 手写 CUDA 可能不如官方 cuDNN/cuBLAS 更快 。
✅ 重写 PyTorch 内核成本高，而且很难长期维护，通常只适用于底层优化团队。
💡 更推荐的方式
-
如果你要写自定义算子，不用改 PyTorch C++，可以用 Python +
torch.autograd.Function：pythonclass MyFunc(torch.autograd.Function): @staticmethod def forward(ctx, x): return x ** 2 @staticmethod def backward(ctx, grad_output): return 2 * grad_output这样就能定义 自定义梯度 ，不用去改
autograd源码。
总结
| 场景 | 是否需要改 PyTorch C++ 底层？ | 推荐方案 |
|---|---|---|
| 训练常见模型（ResNet, Transformer） | ❌ 不需要 | 用 PyTorch API |
| 设计新神经网络结构（NAS） | ❌ 不需要 | 自定义 torch.nn.Module |
| 研究新的损失函数 | ❌ 不需要 | 用 Python 方式写 loss |
| 自定义算子（特殊运算） | ⚠️ 可以改 | 用 torch.autograd.Function 或 C++ 扩展 |
| 分布式训练优化（DDP, NCCL） | ✅ 需要 | 修改 PyTorch DDP 代码 |
| GPU 计算优化（CUDA Kernel） | ✅ 需要 | 编写新的 CUDA 核心算子 |
| 自动微分优化（新的反向传播算法） | ✅ 需要 | 修改 autograd 代码 |
| 硬件支持（TPU, FPGA, AI 芯片） | ✅ 需要 | 增加新硬件后端 |
结论
🚀 一般做深度学习算法（训练/调参/改模型），不用改 PyTorch C++ 层，Python API 足够。
⚡ 如果你是做深度学习底层优化（CUDA、分布式、TPU），那可能需要修改 PyTorch C++ 内核。
如果你的目标是写更高效的深度学习模型，推荐 先深入学习 PyTorch Python API，只有当你需要真正优化计算效率时，才去研究 PyTorch C++ 层！💡
PyTorch 动态计算图与自动求导底层实现
PyTorch 采用**动态图计算（Dynamic Computational Graph, DCG）**的方式，每次前向传播都会构建一张新的计算图，记录所有计算操作，并在反向传播时沿着计算图自动计算梯度。相比 TensorFlow 1.0 的静态计算图，PyTorch 的动态图更直观、易调试，适用于研究和生产。
接下来，我会从 原理 → C++ 实现 → Python API → 实战代码 ，一层一层讲清楚 PyTorch 的自动微分系统。
1. 计算图（Computation Graph）
PyTorch 的计算图是 基于有向无环图（DAG, Directed Acyclic Graph） 的，其中：
- 节点（Node） 代表一个操作（如加法、乘法）。
- 边（Edge） 代表数据流（即 Tensor 之间的依赖关系）。
- 叶子节点（Leaf Node） 是输入变量（如
x = torch.tensor(2.0, requires_grad=True)）。 - 根节点（Root Node） 是最终计算结果（如
z = f(x)）。
计算图示例
python
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 # y = x^2
z = y + 3 # z = x^2 + 3
z.backward() # 计算 dz/dx
print(x.grad) # 输出 4计算图
x
\
^2 (PowBackward)
\
+3 (AddBackward)
|
dz/dx计算 z.backward() 时：
AddBackward计算dz/dy = 1。PowBackward计算dy/dx = 2x。- 结合链式法则
dz/dx = dz/dy * dy/dx = 1 * 2x，所以x.grad = 4。
2. 计算图的 C++ 底层实现
PyTorch 的计算图由 Node 和 Edge 组成，底层的 Tensor 通过 grad_fn 记录其计算来源。
2.1 Tensor 结构
cpp
struct TensorImpl {
caffe2::TypeMeta dtype_; // 数据类型
at::Storage storage_; // 存储数据
at::Device device_; // 设备（CPU / CUDA）
std::shared_ptr<Node> grad_fn_; // 计算图中的前向操作
bool requires_grad_; // 是否需要梯度
};每个 Tensor 通过 grad_fn_ 指向它的计算节点。
2.2 Node 结构
cpp
struct Node {
virtual void apply() = 0; // 计算梯度
std::vector<Tensor> next_edges_; // 指向下一个计算节点
};每个 Node 代表一个计算操作，例如 MulBackward 计算乘法的梯度，AddBackward 计算加法的梯度。
2.3 计算图的构建
在 torch/csrc/autograd/variable.cpp 中，当执行 y = x ** 2 时：
cpp
Tensor pow(const Tensor& self, double exponent) {
auto result = at::empty(self.sizes(), self.options());
at::native::pow_out(result, self, exponent);
if (self.requires_grad()) {
auto grad_fn = std::make_shared<PowBackward>(self, exponent);
result.set_grad_fn(grad_fn);
}
return result;
}- 创建
PowBackward计算节点。 result.set_grad_fn(grad_fn)记录y的梯度来源。- 计算
z = y + 3时，继续构建AddBackward计算节点。
3. 反向传播（Backpropagation）
PyTorch 采用 反向模式自动微分（Reverse-mode Autograd），沿计算图逆向传播梯度。
3.1 反向传播的核心 C++ 代码
在 torch/csrc/autograd/engine.cpp 中，PyTorch 采用 后进先出（LIFO） 调度梯度计算：
cpp
void Engine::execute(Node* graph_root) {
std::deque<Node*> queue;
queue.push_back(graph_root);
while (!queue.empty()) {
Node* node = queue.back(); // 取出 LIFO 结构的最后一个节点
queue.pop_back();
node->apply(); // 计算该节点的梯度
for (Edge& edge : node->next_edges_) {
queue.push_back(edge.function);
}
}
}梯度计算
cpp
struct PowBackward : public Node {
Tensor input_;
double exponent_;
void apply() override {
Tensor grad_input = exponent_ * input_.pow(exponent_ - 1);
input_.grad().add_(grad_input);
}
};执行 z.backward() 时：
AddBackward计算dz/dy = 1PowBackward计算dy/dx = 2x- 计算
dz/dx = dz/dy * dy/dx，并存入x.grad
4. Python API 的实现
PyTorch 通过 torch.autograd.Function 提供 自定义梯度 机制：
python
class MyPow(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x) # 存储 x 以备反向传播使用
return x ** 2
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
return grad_output * 2 * x # 计算梯度使用：
python
x = torch.tensor(2.0, requires_grad=True)
y = MyPow.apply(x)
y.backward()
print(x.grad) # 输出 45. 综合案例
结合所有 PyTorch 知识点，包括：
- 计算图 (
torch.autograd) - 自定义梯度 (
Function) - GPU 加速 (
CUDA) - 动态计算图 (
requires_grad) - 反向传播 (
backward) - 优化器 (
torch.optim) - 深度学习模型 (
torch.nn)
python
import torch
import torch.nn as nn
import torch.optim as optim
# 1. 定义模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc1 = nn.Linear(2, 4)
self.fc2 = nn.Linear(4, 1)
def forward(self, x):
return self.fc2(torch.relu(self.fc1(x)))
# 2. 训练流程
device = "cuda" if torch.cuda.is_available() else "cpu"
model = MyModel().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()
# 3. 训练数据
x = torch.randn(100, 2).to(device)
y = torch.randn(100, 1).to(device)
for epoch in range(100):
optimizer.zero_grad()
output = model(x)
loss = criterion(output, y)
loss.backward() # 自动计算梯度
optimizer.step()
print("训练完成，最终损失：", loss.item())知识点覆盖
✅ 计算图 (backward)
✅ 自动微分 (requires_grad)
✅ GPU 加速 (to(device))
✅ 神经网络 (torch.nn)
✅ 优化器 (torch.optim)
总结
PyTorch 的自动求导系统依赖于 动态图计算 + 计算图结构 + 反向模式自动微分 ，并通过 C++ autograd 计算梯度，在 Python 层提供了易用的 torch.autograd API，使得梯度计算高效灵活。
综合案例学习
这里是一个完整的 PyTorch 综合案例，涵盖 PyTorch 主要知识点，包括：
✅ 数据处理（Dataset & DataLoader）
✅ 模型构建（torch.nn.Module）
✅ 前向传播（forward）
✅ 自动微分（requires_grad & autograd）
✅ 计算图（动态计算图）
✅ GPU 加速（cuda）
✅ 自定义算子（torch.autograd.Function）
✅ 优化器（torch.optim）
✅ 训练 & 评估
✅ 分布式训练（torch.nn.DataParallel）
✅ 自定义激活函数 & 初始化
✅ 模型保存 & 加载
✅ 可视化（matplotlib & torchvision.utils）
📌 任务：手写数字识别（MNIST），使用卷积神经网络（CNN）
- 使用
torchvision.datasets.MNIST处理数据 - 使用
torch.nn.Conv2d构建 CNN - 自定义 Swish 激活函数
- 实现
backward()自定义梯度计算 - 在 GPU 上训练
- 保存 & 载入模型
- 使用
matplotlib可视化训练结果
📜 代码开始
python
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.autograd as autograd
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
import matplotlib.pyplot as plt
import numpy as np
# ========================== 1. 设备配置 ==========================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# ========================== 2. 处理数据 ==========================
# 数据预处理：转换为 Tensor 并归一化
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
# 下载 MNIST 数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transform, download=True)
# DataLoader 处理批次数据
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)
# ========================== 3. 自定义 Swish 激活函数 ==========================
class Swish(autograd.Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return input * torch.sigmoid(input)
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
sigmoid_x = torch.sigmoid(input)
return grad_output * (sigmoid_x + input * sigmoid_x * (1 - sigmoid_x))
# ========================== 4. 定义 CNN 模型 ==========================
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.fc1 = nn.Linear(64 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 10)
# 参数初始化
self._initialize_weights()
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2)
x = x.view(x.size(0), -1) # Flatten
x = Swish.apply(self.fc1(x)) # 使用自定义 Swish 激活函数
x = self.fc2(x)
return x
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
# ========================== 5. 训练函数 ==========================
def train(model, train_loader, criterion, optimizer, epochs=5):
model.train()
for epoch in range(epochs):
total_loss = 0
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f"Epoch [{epoch+1}/{epochs}], Loss: {total_loss / len(train_loader):.4f}")
# ========================== 6. 评估函数 ==========================
def evaluate(model, test_loader):
model.eval()
correct = 0
total = 0
with torch.no_grad(): # 关闭梯度计算，提升推理速度
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f"Test Accuracy: {100 * correct / total:.2f}%")
# ========================== 7. 训练与测试 ==========================
model = CNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
train(model, train_loader, criterion, optimizer, epochs=5)
evaluate(model, test_loader)
# ========================== 8. 保存 & 载入模型 ==========================
torch.save(model.state_dict(), "mnist_cnn.pth")
# 重新载入模型并测试
loaded_model = CNN().to(device)
loaded_model.load_state_dict(torch.load("mnist_cnn.pth"))
evaluate(loaded_model, test_loader)
# ========================== 9. 可视化预测结果 ==========================
def visualize_predictions(model, test_loader):
model.eval()
images, labels = next(iter(test_loader))
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)
fig, axes = plt.subplots(3, 3, figsize=(8, 8))
for i, ax in enumerate(axes.flatten()):
img = images[i].cpu().numpy().squeeze()
ax.imshow(img, cmap="gray")
ax.set_title(f"Label: {labels[i].item()}, Pred: {predicted[i].item()}")
ax.axis("off")
plt.show()
visualize_predictions(model, test_loader)📌 代码总结
这段代码涵盖了 PyTorch 绝大部分核心知识点：
- 数据处理（DataLoader, Dataset）
- 神经网络构建（
torch.nn.Module） - 前向传播 & 反向传播（
forward,backward） - 计算图 & 自动微分（
requires_grad,autograd.Function） - GPU 计算（
to(device)） - 优化器 & 训练（
torch.optim） - 模型保存 & 加载（
torch.save&torch.load） - 可视化预测结果
这可以作为一个 完整的 PyTorch 框架学习案例！🚀🔥