pytorch日积月累3：一次讲清 PyTorch 封装的那些层次

当你在 PyTorch 里写代码的时候，有没有遇到过这样的困惑：“为什么同样的功能，在 PyTorch 里有好几种写法？” 比如，你想做一个简单的线性变换，结果发现有 torch.matmul()，torch.nn.functional.linear()，还有 torch.nn.Linear()，到底哪个该用？它们有什么区别呢？

实际上，PyTorch 提供了不同层级的封装，每个层级有它的独特用处，从底层的计算到高层次的网络搭建。今天我就来带你从底层一路逛到神经网络，看清每一层到底是干啥的、啥时候用。

第一层：CUDA 封装

我们从最底层开始，叫做 CUDA 封装。如果你听起来有点陌生，没关系，打个比方，想象你进了一个超高效的工厂，里面有成百上千的机器人。这些机器人只干一件事：算数。比如说，你让它们做个矩阵乘法，它们就干完这一件事，算完立马走人，啥也不管，数据也不存，就像“干活不留痕迹”的工人。

这就是 CUDA 封装的工作方式。它直接跟 GPU 打交道，每次调用都会用 GPU 里的底层算子做计算。比如做个矩阵乘法，就是这么干的：

举个例子：

# 调用底层的 CUDA 算子做矩阵乘法GEMM_cuda.fwd(mat1, mat2)GEMM_cuda.bwd(grad_o, mat1, mat2)  # 反向传播计算梯度

这层封装适合啥场景呢？性能极致优化！如果你追求速度，想让 GPU 飞起来跑，这就是你的地盘。但一般来说，你可能不太会直接接触到这玩意儿，除非你是做 GPU 算子优化的工程师。

第二层：Autograd 封装

接下来，我们往上走一层，来到了 Autograd 封装。这层和前面的 CUDA 封装不同，它不光计算，还会记住中间结果，方便后面用。我们再打个比方，它就像是个记忆力超强的学生，每次计算之前都会在笔记本上记下重要步骤，等到反向传播的时候再把这些笔记翻出来。

比如，你要自定义一个 ReLU 激活函数，Autograd 就帮你记住前面步骤，等反向传播的时候，用保存的信息帮你计算梯度。

来看代码：

import torchclass TensorModelParallelCrossEntropy(torch.autograd.Function):    @staticmethod    def forward(ctx, logits, target, label_smoothing=0.0):        # do something        ctx.save_for_backward(...)        return loss
    @staticmethod    def backward(ctx, grad_output):        # do something        ... = ctx.saved_tensors        return grad_input, None, None

这种算子可以通过下面这种方式调用：

ce_loss = TensorModelParallelCrossEntropy.apply(logits,labels)

你可以看到调用这种算子并不是通过使用它的forward或者backward函数，而是使用apply函数。

这里torch会进行一些封装，例如调用apply后，会用forward进行计算，并将backward添加到tensor的grad_fn属性计算图中，求导时自动调用。会在使用了torch.no_grad()上下文，不需要求导时，自动抛弃掉save_for_backward存储的张量。

但是这种层级用的也不是太常见，首先观察forward函数的输入参数和backward的输出参数。backward函数返回的梯度数量必须和forward输入参数的数量相同，但是可以用None占位。比如target是标签，label_smoothing是超参，不可学习，不需要导数，这里就会用None占位。因此当你需要某一个功能的时候，需要严格的选择你需要的autograd算子，达到最佳的计算效率，不需要计算的东西不要算。

第三层：Function 封装

再往上一层，就是 Function 封装，这层封装让代码更加灵活。它不仅能做你想要的计算，还能帮你检查输入是否合理、填补一些默认值、应对各种情况。比如，你调用 torch.nn.functional.linear() 这个函数时，它会自动帮你检测输入、输出的维度是否匹配、处理默认的 bias，你几乎不需要管其他的细节。

看看它的使用：

from torch.nn.functional import linear, dropout
# 线性变换和dropout操作output = linear(input, weight, bias)output = dropout(output, p=0.5, training=True)

Function 封装对于大多数常见的操作已经足够了，很多人用 PyTorch 训练模型时，就是直接用这一级别的封装。它能帮你省心很多，不需要你关心太多复杂的细节。

适合场景：绝大多数日常任务。这层封装已经足够灵活和健壮，特别适合直接拿来就用。

第四层：Module 封装

最后，我们来到了最顶层：Module 封装。这层就厉害了，是用来构建整个神经网络的，它不仅帮你管理计算，还能自动保存和初始化模型的参数，比如权重和偏置。

看看怎么用 Module 封装定义一个线性层：

import torch.nn as nn
class MyLinearLayer(nn.Module):    def __init__(self, in_features, out_features):        super(MyLinearLayer, self).__init__()        self.weight = nn.Parameter(torch.randn(out_features, in_features))        self.bias = nn.Parameter(torch.randn(out_features))
    def forward(self, input):        return torch.matmul(input, self.weight.T) + self.bias

你只需要专注于网络结构的设计，剩下的权重初始化、参数管理等，都交给 Module 来搞定。当你用它来搭建网络时，感觉就像是在用乐高搭积木——每块积木（网络层）都帮你准备好，你只要组合就行了。

适合场景：搭建整个神经网络！当你需要全套服务，包括参数管理、前向和反向传播的自动处理，Module 就是你最好的选择。

总结：到底用哪个封装？

最后，来帮你梳理一下到底该用哪个封装：

CUDA 封装：底层操作，适合追求极致性能的人。

Autograd 封装：想自定义正向和反向传播？这个是你的首选！

Function 封装：日常使用，功能强大且灵活。

Module 封装：构建复杂的神经网络模型，省心又方便。

如果你只是想快速搭建一个神经网络，Module 封装就够了。如果你要实现一些自定义的激活函数或层，那你可能要用到 Autograd 封装。而如果你想深入 GPU 优化的世界，搞些硬核的操作，那 CUDA 封装 就是你要玩的地方。