想真正搞懂 PyTorch 的backward()背后发生了什么最好的办法就是读一个麻雀虽小、五脏俱全的框架源码。今天给大家介绍一个用NumPy CuPy从零实现张量计算与自动微分的神经网络框架——Riemann。GitHub 仓库https://github.com/xiangfei2017/RiemannGitee 仓库https://gitee.com/xfcode2021/Riemann在线文档英文Riemann Documentation在线文档中文Riemann 使用指南文档一、Riemann 是什么Riemann是一个类似于 PyTorch 的神经网络编程框架支持张量计算的自动微分提供搭建神经网络所需的全套组件定位是学习、教学与科研。它的接口设计刻意与 PyTorch 保持一致会用 PyTorch 的同学可以零成本上手同时源码足够精简核心张量类就在一个文件里非常适合用来拆开看深度学习框架的内部原理。1.1 核心能力一览张量计算与自动微分支持 0 维到多维张量运算、复数张量、CUDA 加速支持反向传播自动求梯度、函数微分jacobian/hessian、自定义梯度函数线性代数批量矩阵乘、矩阵分解、特征值、行列式、求逆/伪逆、解线性方程组、范数与秩等张量操作reshape、扩缩维、索引切片、gather/scatter、拼接/拆分等神经网络模块Linear、Dropout、BatchNorm、LayerNorm、Embedding、Conv1d/2d/3d、池化、KAN 网络、TransformerMultiheadAttention、TransformerEncoder/Decoder等损失与激活MSE、CrossEntropy、ReLU、Sigmoid、Softmax……优化器SGD、Adam、Adagrad、LBFGS 及学习率调度器视觉模块接口对标 torchvision 的 MNIST/CIFAR10 数据集与 transforms模型序列化张量/模型的存取1.2 模块结构riemann包的文件目录结构 riemann/ ├── tensordef.py # 核心Tensor 类定义 全部张量算子建图就发生在这里 ├── autograd/ # 自动微分backward、grad、Function、jacobian、hessian、gradcheck ├── linalg.py # 线性代数函数 ├── nn/ # 神经网络module/conv/activation/loss/transformer/kan/init ├── optim/ # 优化器与学习率调度器 ├── utils/ # Dataset、DataLoader 等数据工具 ├── vision/ # 数据集MNIST/CIFAR10与图像变换 transforms ├── cuda.py # 设备管理CPU/GPU 切换、Device 上下文 ├── dtype.py # 数据类型 ├── gradmode.py # 梯度模式控制no_grad / enable_grad └── serialization.py # 对象保存与加载 riemann包的模块结构 riemann # 主包 ├── autograd # 自动微分模块 │ └── functional # 自动微分函数式接口 ├── linalg # 线性代数模块 ├── nn # 神经网络模块 │ └── functional # 神经网络函数 ├── optim # 优化器模块 │ └── lr_scheduler # 学习率调度器模块 ├── utils # 工具函数模块 │ └── data # 数据处理工具 ├── vision # 计算机视觉模块 │ ├── datasets # 数据集类 │ └── transforms # 图像变换操作 └── cuda # CUDA/GPU支持1.3 与 PyTorch 的接口对照Riemann 的函数/类命名和调用方式与 PyTorch 高度一致功能PyTorchRiemann创建张量torch.tensor(..., requires_gradTrue)riemann.tensor(..., requires_gradTrue)反向传播loss.backward()loss.backward()线性层nn.Linear(784, 10)r.nn.Linear(784, 10)损失函数nn.CrossEntropyLoss()r.nn.CrossEntropyLoss()优化器optim.SGD(model.parameters(), lr0.01, momentum0.9)r.optim.SGD(model.parameters(), lr0.01, momentum0.9)关闭梯度torch.no_grad()r.no_grad()设备迁移model.to(cuda)model.to(cuda)数据加载torch.utils.data.DataLoaderr.utils.DataLoader数据集torchvision.datasets.MNISTr.vision.datasets.MNIST可以说把 PyTorch 代码里的torch换成riemann大部分脚本就能跑。二、张量计算的底层NumPy 管 CPUCuPy 管 GPU2.1 Tensor 就是 ndarray 的一层皮Riemann 的张量类叫TNTensor Numerics定义在 tensordef.py。它本身不做数值计算真正干活的是底层的数组在 CPU 上底层是NumPy的np.ndarray在 GPU 上底层是CuPy的cp.ndarrayCuPy 可以理解为GPU 版 NumPy——它的 APIcp.asarray、cp.sum、cp.matmul……与 NumPy 几乎一模一样。正因为如此Riemann 的算子代码写一遍就能同时在 CPU 和 GPU 上运行不需要维护两套实现。2.2 CuPy 是可选的导入失败就回退 CPUcuda.py 的开头是典型的优雅降级try: import cupy as cp try: device_count cp.cuda.runtime.getDeviceCount() if device_count 0: raise RuntimeError(No CUDA devices found) test_array cp.full((10, 10), 5.0) # 真正跑一次 CUDA 操作验证 test_result test_array.sum() CUPY_AVAILABLE True except Exception as e: print(fWarning: CuPy imported but CUDA is not functional ({e})...) CUPY_AVAILABLE False cp None except ImportError: CUPY_AVAILABLE False cp None装了 CuPy 且 CUDA 可用 → GPU 模式没装或驱动不对 →cp None框架自动只用 CPU对上层完全透明。2.3 张量住在哪块设备上看 data 的类型就知道TN 的device属性是动态推断的tensordef.py:787-806property def device(self): if cp and isinstance(self.data, cp.ndarray): device_idx self.data.device.id return Device(fcuda:{device_idx}) # GPU 张量 else: return Device(cpu) # CPU 张量创建张量时tensor()工厂函数根据device参数决定把数据放进 NumPy 还是 CuPytensordef.py:4896-4925if use_cuda: if isinstance(data, np.ndarray): with cp.cuda.Device(target_device_idx): arr cp.asarray(data) # CPU - GPU ... else: if isinstance(data, np.ndarray): arr data # 本来就是 CPU else: arr cp.asnumpy(data) # GPU - CPU2.4.to()设备迁移CPU↔GPU 的数据搬运在TN.to()中完成tensordef.py:932-964核心就是两个函数CPU → GPUcp.asarray(data)还支持non_blockingTrue走异步流GPU → CPUcp.asnumpy(data)所以用起来和 PyTorch 一模一样import riemann as r device r.device(cuda if r.cuda.is_available() else cpu) x r.tensor([1.0, 2.0, 3.0]).to(device) # 张量上 GPU model.to(device) # 整个模型上 GPU小结NumPy 提供 CPU 张量算力CuPy 提供 GPU 张量算力TN 类在外面包一层并记录计算图——这就是 Riemann 全部的秘密。三、核心原理反向传播自动求导是怎么实现的这是本文的重点。Riemann 采用和 PyTorch 一样的基于动态计算图的反向模式自动微分reverse-mode autodiff。3.1 计算图长什么样每次前向运算结果张量都会记下我是由谁、通过什么运算算出来的。TN 类上有这几个关键属性tensordef.py:113-123属性含义data底层 np/cp 数组存实际数值fromvars元组当前张量由哪些输入张量计算而来图的边gradfuncs元组与 fromvars 一一对应的梯度钩子函数图的边上怎么传梯度parms元组运算用到的非张量参数如 sum 的 dimrequires_grad是否需要梯度grad叶子节点最终存放梯度的地方grad_value反向传播过程中的临时梯度缓存rcv_grad_count反向传播过程中记录该节点还有多少条路径的梯度没收到计数器什么是梯度钩子函数读者可以把计算图想象成一张水管网络前向是水流数据从源头流向出口反向则是梯度从出口倒流回源头。每条管道图的边上都装了一个钩子——它负责把流到当前节点的梯度按照这个算子的数学导数规则转换成应该继续往上游传的梯度。所以叫它梯度钩子函数它钩住梯度、按规则变换、再放行给上游。每个算子加、乘、卷积……都注册了自己专属的梯度钩子函数。以加法为例tensordef.py:2159-2187Python# 运算左值为self右值为TNnumpy数组listtuple整数或浮点数 def __add__(self,right_obj): # 获取self的设备 dev self.device if isinstance(right_obj,TN): if dev ! right_obj.device: raise RuntimeError(fExpected all tensors to be on the same device, but found at least two devices, {dev} and {right_obj.device}!) right_data right_obj.data right_requires_grad right_obj.requires_grad else: right_data right_obj right_requires_grad False #requires_grad属性在运算时传递到结果tensor requires_grad (is_grad_enabled() and (self.requires_grad or right_requires_grad)) ret tensor(self.data right_data, devicedev, requires_gradrequires_grad) if requires_grad: if self.requires_grad and right_requires_grad: ret.fromvars(self,right_obj) ret.gradfuncs(_add_grad_left,_add_grad_right) elif self.requires_grad: ret.fromvars(self,) ret.gradfuncs(_add_grad_left,) else: ret.fromvars (right_obj,) ret.gradfuncs (_add_grad_right,) return ret前向每走一步就挂一个新节点到图上。所谓动态图就是图在运行时一边执行一边搭建不用先编译。3.2 梯度钩子函数链式法则的积木每个算子只需要负责自己的局部导数。比如加法c a b∂c/∂a 1所以上游梯度原样传给 atensordef.py:7429-7441Pythondef _add_grad_left(result_tensor: TN, i: int) - TN: left_tensor result_tensor.fromvars[i] left_var_shape left_tensor.shape result_shape result_tensor.shape # shape一样时直接返回grad_value否则对result_tensor.grad_value进行sum缩减 if left_var_shape result_shape: grad result_tensor.grad_value else: # left_tensor与result_tensor的shape比较获取需left_tensor广播轴序号的元组 broadcast_axes _get_broadcast_axis(result_shape, left_var_shape) grad sum(result_tensor.grad_value, dimbroadcast_axes, keepdimFalse)._reshape(left_var_shape) return grad def _add_grad_right(result_tensor:TN, i:int)-TN: right_tensor result_tensor.fromvars[i] right_var_shape right_tensor.shape result_shape result_tensor.shape # shape一样时直接返回grad_value否则对result_tensor.grad_value进行sum缩减 if right_var_shape result_shape: grad result_tensor.grad_value else: # right_tensor与result_tensor的shape比较获取需right_tensor广播轴序号的元组 broadcast_axes_get_broadcast_axis(result_shape,right_var_shape) grad sum(result_tensor.grad_value,dimbroadcast_axes,keepdimFalse)._reshape(right_var_shape) return grad如果加法发生过广播比如(3,) (2,3)梯度需要把广播出去的轴 sum 回来再 reshape 回原始形状——这正是上面else分支在做的事。乘法c a * b∂c/∂a b所以传给 a 的梯度 上游梯度 × btensordef.py:7464-7475Pythondef _mul_grad_left(result_tensor:TN, i:int)-TN: left_tensor result_tensor.fromvars[i] # parms为空时从fromvars中获取右值否则从parms中获取 if result_tensor.parms: right_value result_tensor.parms[i] else: right_value result_tensor.fromvars[i1] left_var_shape left_tensor.shape result_shape result_tensor.shape # 右值可能是非张量使用conj()函数获取共轭 left_grad result_tensor.grad_value * conj(right_value) # shape一样时直接返回left_grad否则对left_grad进行sum缩减 if left_var_shape result_shape: grad left_grad else: # left_tensor与result_tensor的shape比较获取需left_tensor广播轴序号的元组 broadcast_axes_get_broadcast_axis(result_shape,left_var_shape) grad sum(left_grad, dimbroadcast_axes, keepdimFalse)._reshape(left_var_shape) return grad def _mul_grad_right(result_tensor:TN, i:int)-TN: # parms为空时从fromvars中获取左值否则从parms中获取 if result_tensor.parms: left_tensor result_tensor.parms[i] else: left_tensor result_tensor.fromvars[i-1] right_tensor result_tensor.fromvars[i] right_var_shape right_tensor.shape result_shape result_tensor.shape # 左值一定是张量使用conj()成员函数获取共轭效率更高 right_grad result_tensor.grad_value * left_tensor.conj() # shape一样时直接返回right_grad否则对right_grad进行sum缩减 if right_var_shape result_shape: grad right_grad else: # right_tensor与result_tensor的shape比较获取需right_tensor广播轴序号的元组 broadcast_axes_get_broadcast_axis(result_shape,right_var_shape) grad sum(right_grad, dimbroadcast_axes, keepdimFalse)._reshape(right_var_shape) return grad这里用conj()取共轭是为了同时兼容复数张量——复数求导链式法则里会出现共轭。每个算子只管自己这一环反向传播时再把它们按链式法则串起来。3.3 backward() 的完整流程入口是TN.backward()tensordef.py:3915-3992。整个反向传播分为三大步下面给出完整代码并逐行解读。第 1 步参数校验 播种初始梯度Pythondef backward(self, gradient: TN | None None, retain_graph: bool False, create_graph: bool False): # ---- ① 校验只有 requires_gradTrue 的张量才能反向传播 ---- if not self.requires_grad: raise RuntimeError(Only a tensor require grad can call backward()) if gradient is None: # ---- ② 没传 gradient要求 self 必须是标量 ---- if self.data.ndim 0: raise RuntimeError(Only a scalar can call backward() without argument grad_outputs) if self.is_complex(): raise RuntimeError(fgrad can be implicitly created only for real scalar outputs but got {self.dtype}) # ---- ③ 初始化计算图清缓存 统计入边数---- need_hook_processing self._init_calc_graph() # ---- ④ 播种标量对自身的导数 1.0 ---- self.grad_value tensor(1.0, dtypeself.dtype, deviceself.device, requires_gradcreate_graph) elif isinstance(gradient, TN): # ---- ② 用户传了外部 gradient形状必须匹配 ---- if gradient.data.shape self.data.shape: need_hook_processing self._init_calc_graph() self.grad_value gradient.to(self.device).detach().requires_grad_(create_graph) else: raise RuntimeError(shape of gradient need to be same as the shape of outputs) else: raise TypeError(fgradient can be either tensor or None, but got {type(gradient)}) # ---- ⑤ 真正执行反向传播 ---- with set_grad_enabled(create_graph): if need_hook_processing: self._backward_with_hooks() # 带模块钩子的慢路径 else: self._backward_without_hooks() # 无钩子的快速路径 return逐行解读①backward()只能在需要梯度的张量上调用。如果你忘了requires_gradTrue这里直接报错。②如果用户没传gradientRiemann 要求self必须是标量ndim 0。这是深度学习最常见场景——loss是一个标量loss.backward()隐式使用梯度1.0。如果你对一个非标量张量调用backward()又不传 gradient会报错这和 PyTorch 行为一致。复数标量也不行因为复数没法隐式播种梯度。③_init_calc_graph()是准备工作下一小节详述。它返回一个布尔值need_hook_processing告诉你计算图里是否存在注册了反向钩子的nn.Module——这决定了走快路径还是慢路径。④标量对自身的导数恒为 1所以根节点self.grad_value 1.0。注意这个 1.0 也是一个TN张量和self同设备同 dtype如果create_graphTrue它也参与建图用于高阶导数。②如果用户传了gradient比如向量-Jacobian 积场景形状必须和self一致梯度会被 detach断开计算图再搬到self所在设备。⑤with set_grad_enabled(create_graph)控制梯度追踪开关create_graphTrue时反向传播中产生的梯度运算本身也会被记录到新计算图里从而可以对梯度再求梯度算二阶导。然后根据有没有模块钩子分两条路径执行。第 2 步_init_calc_graph()—— 清缓存 统计入边数反向传播前必须先知道每个节点会被多少条边喂梯度。为什么因为一个张量可能被多个下游节点使用比如y a*a a变量a有两条路径回传梯度。必须等所有路径的梯度到齐、累加完才能继续往它的上游传否则梯度不完整。Pythondef _init_calc_graph(self) - bool: stack [self] # 从根节点loss出发 visited {self} # visited 集合防止重复访问图里可能有菱形依赖 self.rcv_grad_count 0 # 根节点入边计数清零 need_hook_processing False while stack: item: TN stack.pop(-1) # DFS 深度优先 item.grad_value None # ★ 把所有节点的临时梯度缓存清空 # 检查是否需要模块钩子处理短路逻辑一旦确定为True就不再检查 if not need_hook_processing and item._output_of is not None: module item._output_of if module._backward_pre_hooks or module._backward_hooks: need_hook_processing True # 叶子节点没有 fromvars不再往下走 if item.is_leaf: continue # 遍历当前节点的所有来源上游张量 varlist item.fromvars for var in varlist: if var.requires_grad: if var not in visited: var.rcv_grad_count 0 # 首次访问清零 stack.append(var) visited.add(var) var.rcv_grad_count 1 # ★ 每有一条入边计数 1 return need_hook_processing逐行解读stackvisited标准 DFS。visited用集合set而非列表查找 O(1)因为计算图里可能存在菱形依赖一个节点被多条路径引用不防重复就会把rcv_grad_count算错。item.grad_value None把上一次反向传播可能残留的临时梯度全部清空保证干净起步。need_hook_processing检查item._output_of记录当前张量是哪个 nn.Module 的输出。如果那个 Module 注册了_backward_pre_hooks或_backward_hooks就需要走慢路径。这里用短路逻辑——一旦确定为 True后面所有节点都不再检查节省时间。if item.is_leaf: continue叶子节点用户创建的参数/输入fromvars为空是图的末端没有上游可传跳过。var.rcv_grad_count 1核心。遍历item.fromvars每遇到一个需要梯度的上游变量它的入边计数 1。最终rcv_grad_count的值就是这个节点在反向传播中应该收到几次梯度。等收到次数等于这个值节点才梯度就绪、可以入栈继续往上传播。第 3 步_backward_without_hooks()—— 栈式反向传播主循环这是最常见路径没有模块钩子时走这里也是整个自动求导的心脏tensordef.py:3892-3913Pythondef _backward_without_hooks(self) - None: stack [self] # 栈里放梯度已就绪的节点初始只有 loss while stack: item: TN stack.pop(-1) # ① 弹出一个梯度已就绪的节点 if item._save_grad(): # ② 保存梯度若是叶子节点则返回 True continue # 叶子节点无需再向上传直接进入下一轮 vars_received_grad set() item._propagate_grad_to_sources(vars_received_grad) # ③ 把梯度传给上游 item.grad_value None # ④ 传完即清释放内存 stack.extend(vars_received_grad) # ⑤ 新就绪的节点入栈继续往上逐行解读①stack.pop(-1)从栈尾弹出O(1)本质是 DFS。栈里永远只放梯度已收集完毕的节点——这是这个算法最关键的不变式invariant。②_save_grad()做两件事把grad_value写进.grad叶子节点总是写中间节点仅当retains_gradTrue才写并返回是否是叶子节点。叶子节点fromvars为空没有上游continue跳过后续传播。③_propagate_grad_to_sources()是传播的核心下一小节详述。④item.grad_value None节点传完梯度使命结束立即清空临时缓存释放内存。这一点对大模型很关键——否则反向传播过程中所有中间梯度都留着显存会爆。⑤stack.extend(vars_received_grad)把刚凑齐梯度的上游节点们入栈下一轮循环处理它们。如此反复直到栈空——所有叶子节点的.grad都算完了。Pythondef _save_grad(self:TN): 将self的grad_value保存到self的grad中返回是否是叶子节点 1._backward_with_hooks或_backward_without_hooks里调用_save_grad self是出栈节点已保证了self是一个requires_gradTrue的叶子节点或中间节点。 2.叶子节点is_leafTrue总是保存梯度到self.grad 中间节点只有在retains_gradTrue时才保存梯度。 is_leaf not self.fromvars # 叶子节点总是保存梯度 # 注意is_leaf为True意味着requires_gradTrue且没有fromvars if is_leaf: if self.grad is None: # 如还没有梯度值直接赋值 self.grad self.grad_value else: # 如已有梯度值累计梯度注意不要使用原地 self.grad self.grad self.grad_value # 叶子节点存完梯度后梯度缓存置None self.grad_value None # 中间节点只有retains_gradTrue时才保存梯度 elif self.retains_grad: if self.grad is None: self.grad self.grad_value else: self.grad self.grad self.grad_value # 中间节点保存梯度后不清空grad_value因为还要传播 return is_leaf第 3 步续_propagate_grad_to_sources()—— 调用梯度钩子函数把梯度传给上游Pythondef _propagate_grad_to_sources(self: TN, vars_received_grad: set) - None: fromvars self.fromvars # 当前节点的来源张量列表 gradfuncs self.gradfuncs # 与 fromvars 一一对应的梯度钩子函数列表 for i in range(len(fromvars)): var: TN fromvars[i] # 第 i 个上游张量 fn gradfuncs[i] # 对应的梯度钩子函数 if var.requires_grad: # ① 调用梯度钩子函数算出应该回传给 var 的梯度 rcv_grad_value: TN fn(self, i) # ② 累加到 var.grad_value不是覆盖 var._add_received_grad_value(rcv_grad_value) # ③ 计数器 -1 var.rcv_grad_count - 1 # ④ 计数器归零所有入边的梯度都到齐了var 可以入栈继续向上传 if var.rcv_grad_count 0: vars_received_grad.add(var)逐行解读①fn(self, i)这是反向传播的灵魂一行。fn就是前向时登记的梯度钩子函数比如_add_grad_left、_mul_grad_left。它接收self当前节点携带已就绪的grad_value和索引i内部按算子的数学导数规则把self.grad_value变换成应该传给第i个上游的梯度。这正是链式法则的一环∂L/∂var ∂L/∂self × ∂self/∂var而fn内部封装的就是∂self/∂var。②_add_received_grad_value()把算出的梯度累加到var.grad_value。为什么是累加因为var可能被多个下游引用如y a*a aa有两条回传路径。看它的实现tensordef.py:3798-3823Pythondef _add_received_grad_value(self: TN, grad_value: TN): if grad_value is None: return # dtype 对齐防止 float32 被提升成 float64 if grad_value.dtype ! self.dtype: warnings.warn(...) grad_value grad_value.type(self.dtype) if self.grad_value is None: self.grad_value grad_value # 第一次直接赋值 else: # ★ 注意不能用原地 self.grad_value self.grad_value grad_value源码注释特意强调了不能原地像c a b这种情况c 的grad_value会原封不动传给 a 和 b两者底层共享同一个数组如果原地加a 的累加会污染 b 的梯度。所以这里用self.grad_value grad_value创建新数组。③var.rcv_grad_count - 1每收到一条边的梯度计数减一。④if var.rcv_grad_count 0计数归零意味着所有下游喂过来的梯度都已到齐并累加完毕这个节点现在梯度就绪加入vars_received_grad集合等主循环把它入栈。带模块钩子的慢路径_backward_with_hooks()如果计算图里有nn.Module注册了反向钩子用于自定义模块级的梯度处理比如梯度裁剪就走这条路tensordef.py:3853-3890。结构和快路径几乎一样只是多了一层BackwardHookManagerPythondef _backward_with_hooks(self) - None: hook_manager BackwardHookManager() stack [self] while stack: item: TN stack.pop(-1) # 在反向传播前统一检查和处理两种钩子 items_for_propagate hook_manager.process_module_hooks(item) # 在调用钩子之后保存梯度确保钩子可以访问所有输入的梯度grad_value item._save_grad() # 如果没有可传播节点模块输入或输出梯度未收集齐跳过本次循环 if not items_for_propagate: continue vars_received_grad set() for current_item in items_for_propagate: current_item._propagate_grad_to_sources(vars_received_grad) current_item.grad_value None # 传播完清空 stack.extend(vars_received_grad) hook_manager.cleanup()区别在于hook_manager.process_module_hooks(item)会在节点出栈时检查反向预处理钩子full_backward_pre_hook模块所有输出梯度就绪时调用可修改输出梯度反向钩子full_backward_hook模块所有输入梯度就绪时调用可修改输入梯度如果没有注册钩子第 1 步的_init_calc_graph()就返回False直接走快路径避免任何额外开销。一张图总结整个反向传播流程前向建图: x ──┐ ├─(mul)── y ──(add)── loss (标量) w ──┘ ▲ b ───────────┘ backward() 执行流程: 第1步 loss.grad_value 1.0 ← 播种 第2步 _init_calc_graph() ← DFS 统计每个节点的 rcv_grad_count 例x.rcv_grad_count1, w.rcv_grad_count1, b.rcv_grad_count1 第3步 栈[loss] 弹出 loss → _save_grad(非叶子,跳过) → _propagate_grad_to_sources: 调用 add 的梯度钩子 → y.grad_value 1, b.grad_value 1 y.rcv_grad_count 1→0 ✓入栈 b.rcv_grad_count 1→0 ✓入栈 loss.grad_value None 栈[y, b] 弹出 b → _save_grad(叶子) → .grad 1 ← 叶子梯度写盘continue 弹出 y → _save_grad(非叶子,跳过) → _propagate_grad_to_sources: 调用 mul 的梯度钩子: dx grad_y * w → x.grad_value dx x.rcv_grad_count 1→0 ✓入栈 dw grad_y * x → w.grad_value dw w.rcv_grad_count 1→0 ✓入栈 y.grad_value None 栈[x, w] 弹出 w → _save_grad(叶子) → .grad dw 弹出 x → _save_grad(叶子) → .grad dx 栈[] → 结束 最终x.grad、w.grad 就绪 → 优化器用它们更新参数3.4 还有这些高级能力grad(outputs, inputs)不修改.grad直接返回梯度张量而且只算到指定 inputs 就提前停止算 jacobian/hessian 时更省内存autograd/grad.py:99-271高阶导数backward(create_graphTrue)或grad(..., create_graphTrue)让求梯度这件事本身也建图从而可以二阶梯度下降、算 Hessian自定义求导继承autograd.Function实现forward/backward静态方法用ctx.save_for_backward()缓存前向变量MyFunction.apply(x)调用——和 PyTorch 的torch.autograd.Function用法一模一样autograd/grad.py:619-734gradcheck有限差分数值梯度 vs 解析梯度对比方便验证自己写的算子对不对值得一提的是整套反向传播代码里操作的都是TN对象底层是 NumPy 还是 CuPy 数组无所谓——自动求导在 CPU 和 GPU 上行为完全一致。四、实战MNIST 手写数字识别理论说完跑个真实例子。下面代码改编自仓库自带的 examples/nn_MNIST_CE_SGD.py网络结构为784 → 200(ReLU) → 10的全连接网络交叉熵损失 SGD动量 0.9并自动选择 GPUPython Riemann 框架实战MNIST 手写数字识别 网络Flatten - Linear(784,200) - ReLU - Linear(200,10) 损失CrossEntropyLoss 优化器SGD(lr0.01, momentum0.9) import time from tqdm import tqdm import riemann as r from riemann.vision.datasets import MNIST from riemann.vision.transforms import Compose, ToTensor, Normalize # ---------- 1. 自动选择设备 ---------- device r.device(cuda if r.cuda.is_available() else cpu) print(fUsing device: {device}) # ---------- 2. 数据预处理与加载 ---------- transform Compose([ ToTensor(), Normalize((0.1307,), (0.3081,)) # MNIST 官方均值/标准差 ]) train_dataset MNIST(root./data, trainTrue, transformtransform, downloadTrue) test_dataset MNIST(root./data, trainFalse, transformtransform, downloadTrue) train_loader r.utils.DataLoader(train_dataset, batch_size100, shuffleTrue) test_loader r.utils.DataLoader(test_dataset, batch_size100, shuffleFalse) # ---------- 3. 定义模型 ---------- class Classifier(r.nn.Module): def __init__(self): super().__init__() self.model r.nn.Sequential( r.nn.Flatten(), # (B,1,28,28) - (B,784) r.nn.Linear(784, 200), r.nn.ReLU(), r.nn.Linear(200, 10) # 10 个数字类别 ) def forward(self, x): return self.model(x) model Classifier().to(device) # 模型搬到 GPU/CPU criterion r.nn.CrossEntropyLoss() optimizer r.optim.SGD(model.parameters(), lr0.01, momentum0.9) # ---------- 4. 训练循环 ---------- epochs 3 for epoch in range(epochs): model.train() total_loss 0.0 bar tqdm(train_loader, descfEpoch {epoch1}/{epochs}) for data, target in bar: data, target data.to(device), target.to(device) # 数据上设备 output model(data) # 前向 loss criterion(output, target) optimizer.zero_grad() # 清梯度 loss.backward() # 反向传播就是第三章讲的那套流程 optimizer.step() # 参数更新 total_loss loss.item() bar.set_postfix({Loss: f{loss.item():.4f}}) print(fEpoch {epoch1} 平均损失: {total_loss/len(train_loader):.4f}) # ---------- 5. 测试集评估 ---------- model.eval() correct, total 0, 0 with r.no_grad(): # 推理时关闭建图省内存 for data, target in tqdm(test_loader, descTesting): data, target data.to(device), target.to(device) output model(data) pred output.argmax(dim1) total target.size(0) correct (pred target).sum().item() print(f测试集准确率: {100*correct/total:.2f}% ({correct}/{total}))运行方式Bash运行# 安装GPU 用户再加装对应 CUDA 版本的 cupy如 cupy-cuda12x pip install numpy pillow tqdm pip install cupy-cuda12x # 可选没有则自动用 CPU # 运行仓库 examples 目录下也有现成脚本 python examples/nn_MNIST_CE_SGD.py预期输出3 个 epoch 后测试集准确率约 96%~97%Using device: cuda:0 Epoch 1 平均损失: 0.52xx 测试集准确率: 93.xx% (93xx/10000) Epoch 2 平均损失: 0.28xx 测试集准确率: 95.xx% (95xx/10000) Epoch 3 平均损失: 0.23xx 测试集准确率: 96.xx% (96xx/10000)可以看到整个训练脚本的写法与 PyTorch 几乎逐行对应——model.to(device)、zero_grad() → backward() → step()、no_grad()、argmax(dim1)但每一行背后都是你能直接打开阅读的源码。五、总结Riemann是一个接口对标 PyTorch、面向教学与科研的神经网络框架代码精简、可读性强张量底层NumPy 提供 CPU 算力CuPy 提供 GPU 算力TN类封装 ndarray 并动态推断设备自动求导前向时每个算子把来源张量 梯度钩子函数挂到动态计算图上backward()时播种梯度1.0用栈 计数器从损失节点反向遍历调用各算子的梯度钩子函数按链式法则传播梯度叶子节点梯度累加到.grad支持grad()、高阶导数create_graphTrue、自定义Function、gradcheck想验证学习效果直接跑 MNIST/CIFAR10 示例即可如果你正在学深度学习、好奇loss.backward()到底干了什么强烈建议把仓库 clone 下来从 tensordef.py 的__add__和backward两个函数开始读——相信我读完之后再看 PyTorch 文档会有原来如此的通透感。 项目地址欢迎 StarGitHub 仓库https://github.com/xiangfei2017/RiemannGitee 仓库https://gitee.com/xfcode2021/Riemann