PyTorch 小课堂！梯度计算详解-技术圈

1. 函数的反向传播

2. 计算图的反向传播

3. 数值梯度检查

4. 在自动求导时检测错误产生路径

5. 设置是否需要梯度

6. model.eval() 与 torch.no_grad()

7. 提供 function 级别的统计信息

小伙伴们大家好呀～ PyTorch 源码解读系列又和大家见面啦！在前面的系列文章中，我们介绍了模型训练的原料——数据的预处理，以及 PyTorch 所提供的 nn.Module 结构解析。接下来，我们将为大家介绍 PyTorch 中的 autograd 模块的功能。其中主要涉及 torch/autograd 目录下代码，未涉及底层的 C++ 实现，请大家放心食用。

函数的反向传播

(torch.autograd.function)

我们在构建网络的时候，通常会使用 PyTorch 所提供的 nn.Module（例如 nn.Conv2d, nn.ReLU 等）作为基本单元。而这些 Module 通常包裹着 autograd function，以其作为真正实现的部分。例如 nn.ReLU 实际使用 torch.nn.functional.relu（F.relu）：

from torch.nn import functional as F
class ReLU(Module):    __constants__ = ['inplace']    inplace: bool
    def __init__(self, inplace: bool = False):        super(ReLU, self).__init__()        self.inplace = inplace
    def forward(self, input: Tensor) -> Tensor:        return F.relu(input, inplace=self.inplace)

这里的 F.relu 类型为 function，若再剥开一层，其实际包裹的函数类型为 builtin_function_or_method，这也是真正完成运算的部分，而这些部分通常使用 C++ 实现（如 ATen）。至此我们已经知道，一个模型的运算部分由 autograd functions 组成，这些 autograd functions 内部定义了 forward，backward 方法用以描述前向和梯度反传的过程，组合后可以实现整个模型的前向和梯度反传。因而以 torch.autograd.function 中所定义的 Function 类为基类，我们可以实现自定义的 autograd function。

需要注意的是，所实现的 function 需包含 forward 及 backward 两个方法。以下以 Exp 和 GradCoeff 两个自定义 autograd function 为例：

class Exp(Function):                    # 此层计算e^x
    @staticmethod    def forward(ctx, i):                # 模型前向        result = i.exp()        ctx.save_for_backward(result)   # 保存所需内容，以备backward时使用，所需的结果会被保存在saved_tensors元组中；此处仅能保存tensor类型变量，若其余类型变量（Int等），可直接赋予ctx作为成员变量，也可以达到保存效果        return result
    @staticmethod    def backward(ctx, grad_output):     # 模型梯度反传        result, = ctx.saved_tensors     # 取出forward中保存的result        return grad_output * result     # 计算梯度并返回
# 尝试使用x = torch.tensor([1.], requires_grad=True)  # 需要设置tensor的requires_grad属性为True，才会进行梯度反传ret = Exp.apply(x)                          # 使用apply方法调用自定义autograd functionprint(ret)                                  # tensor([2.7183], grad_fn=)ret.backward()                              # 反传梯度print(x.grad)                               # tensor([2.7183])

Exp 函数的前向很简单，直接调用 tensor 的成员方法 exp 即可。反向时，我们知道：

$\frac{\partial e^x}{\partial x} = e^x$

因此我们直接使用 $e^x$ 乘以 grad_output 即可得梯度。可以发现，自定义的函数 Exp 正确地进行了前向与反向。

同时，我们可以注意到：前向后所得的结果包含了 grad_fn 属性，这一属性指向用于计算其梯度的函数（即 Exp 的 backward 函数）。接下来我们一起看另一个函数 GradCoeff，其功能是反传梯度时乘以一个自定义系数。

class GradCoeff(Function):                  @staticmethod    def forward(ctx, x, coeff):                 # 模型前向        ctx.coeff = coeff                       # 将coeff存为ctx的成员变量        return x.view_as(x)
    @staticmethod    def backward(ctx, grad_output):             # 模型梯度反传        return ctx.coeff * grad_output, None    # backward的输出个数，应与forward的输入个数相同，此处coeff不需要梯度，因此返回None
# 尝试使用x = torch.tensor([2.], requires_grad=True)ret = GradCoeff.apply(x, -0.1)                  # 前向需要同时提供x及coeff，设置coeff为-0.1ret = ret ** 2                          print(ret)                                      # tensor([4.], grad_fn=)ret.backward()  print(x.grad)                                   # tensor([-0.4000])，梯度已乘以相应系数

计算图的反向传播

(torch.autograd.functional)

前面我们描述了单个函数的反向传播，以及如何编写自定义的 autograd function。下面，我们简单介绍 PyTorch 中所提供的计算图反向传播的接口。

在训练过程中，我们通常利用 prediction 和 groundtruth label 来计算 loss（loss 的类型为 Tensor），随后调用 loss.backward() 进行梯度反传。而 Tensor 类的 backward 方法，实际调用的就是 torch.autograd.backward 这一接口。这一 python 接口实现了计算图级的反向传播。

class Tensor(torch._C._TensorBase)
    def backward(self, gradient=None, retain_graph=None, create_graph=False):        relevant_args = (self,)        ...        torch.autograd.backward(self, gradient, retain_graph, create_graph)        # gradient: 形状与tensor一致，可以理解为链式求导的中间结果，若tensor标量，可以省略（默认为1）        # retain_graph: 多次反向传播时梯度累加。反向传播的中间缓存会被清空，为进行多次反向传播需指定retain_graph=True来保存这些缓存。        # create_graph: 为反向传播的过程同样建立计算图，可用于计算二阶导

在 PyTorch 实现中，autograd 会随着用户的操作，记录生成当前 variable 的所有操作，并建立一个有向无环图（DAG）。图中记录了操作 Function，每一个变量在图中的位置可通过其 grad_fn 属性在图中的位置推测得到。

在反向传播过程中，autograd 沿着这个图从当前变量（根节点 F）溯源，可以利用链式求导法则计算所有叶子节点的梯度。每一个前向传播操作的函数都有与之对应的反向传播函数用来计算输入的各个 variable 的梯度，这些函数的函数名通常以 Backward 结尾。我们构建如下图所示的简化的计算图，并以此为例进行简单介绍。

A = torch.tensor(2., requires_grad=True)B = torch.tensor(.5, requires_grad=True)E = torch.tensor(1., requires_grad=True)C = A * BD = C.exp()F = D + Eprint(F)        # tensor(3.7183, grad_fn=) 打印计算结果，可以看到F的grad_fn指向AddBackward，即产生F的运算print([x.is_leaf for x in [A, B, C, D, E, F]])  # [True, True, False, False, True, False] 打印是否为叶节点，由用户创建，且requires_grad设为True的节点为叶节点print([x.grad_fn for x in [F, D, C, A]])    # [, , , None]  每个变量的grad_fn指向产生其算子的backward function，叶节点的grad_fn为空print(F.grad_fn.next_functions) # ((, 0), (, 0)) 由于F = D + E， 因此F.grad_fn.next_functions也存在两项，分别对应于D, E两个变量，每个元组中的第一项对应于相应变量的grad_fn，第二项指示相应变量是产生其op的第几个输出。E作为叶节点，其上没有grad_fn，但有梯度累积函数，即AccumulateGrad（由于反传时多出可能产生梯度，需要进行累加）F.backward(retain_graph=True)   # 进行梯度反传print(A.grad, B.grad, E.grad)   # tensor(1.3591) tensor(5.4366) tensor(1.) 算得每个变量梯度，与求导得到的相符print(C.grad, D.grad)   # None None 为节约空间，梯度反传完成后，中间节点的梯度并不会保留

我们再来看看下面的计算图，并在这个计算图上模拟 autograd 所做的工作：

A = torch.tensor([3.], requires_grad=True)B = torch.tensor([2.], requires_grad=True)C = A ** 2D = B ** 2E = C * DF = D + E
F.manual_grad = torch.tensor(1)                             # 我们用manual_grad表示，在已知计算图结构的情况下，我们模拟autograd过程手动算得的梯度D.manual_grad, E.manual_grad = F.grad_fn(F.manual_grad)     C.manual_grad, tmp2 = E.grad_fn(E.manual_grad)D.manual_grad = D.manual_grad + tmp2                        # 这里我们先完成D上的梯度累加，再进行反传A.manual_grad = C.grad_fn(C.manual_grad)B.manual_grad = D.grad_fn(D.manual_grad)                    # (tensor([24.], grad_fn=), tensor([40.], grad_fn=))

下面，我们编写一个简单的函数，并在这个计算图上进行 autograd，并验证结果是否正确。

# 这一例子仅可用于每个op只产生一个输出的情况，且效率很低（由于对于某一节点，每次未等待所有梯度反传至此节点，就直接将本次反传回的梯度直接反传至叶节点）def autograd(grad_fn, gradient):    auto_grad = {}    queue = [[grad_fn, gradient]]    while queue != []:        item = queue.pop()        gradients = item[0](item[1])        functions = [x[0] for x in item[0].next_functions]            if type(gradients) is not tuple:            gradients = (gradients, )        for grad, func in zip(gradients, functions):                if type(func).__name__ == 'AccumulateGrad':                if hasattr(func.variable, 'auto_grad'):                    func.variable.auto_grad = func.variable.auto_grad + grad                else:                    func.variable.auto_grad = grad            else:                queue.append([func, grad])
A = torch.tensor([3.], requires_grad=True)B = torch.tensor([2.], requires_grad=True)C = A ** 2D = B ** 2E = C * DF = D + E
autograd(F.grad_fn, torch.tensor(1))print(A.auto_grad, B.auto_grad)         # tensor(24., grad_fn=) tensor(40., grad_fn=)
# 这一autograd同样可作用于编写的模型，我们将会看到，它与PyTorch自带的backward产生了同样的结果from torch import nn
class MLP(nn.Module):    def __init__(self):        super().__init__()        self.fc1 = nn.Linear(10, 5)        self.relu = nn.ReLU()        self.fc2 = nn.Linear(5, 2)        self.fc3 = nn.Linear(5, 2)        self.fc4 = nn.Linear(2, 2)
    def forward(self, x):        x = self.fc1(x)        x = self.relu(x)        x1 = self.fc2(x)        x2 = self.fc3(x)        x2 = self.relu(x2)        x2 = self.fc4(x2)        return x1 + x2
x = torch.ones([10], requires_grad=True)mlp = MLP()mlp_state_dict = mlp.state_dict()
# 自定义autogradmlp = MLP()mlp.load_state_dict(mlp_state_dict)y = mlp(x)z = torch.sum(y)autograd(z.grad_fn, torch.tensor(1.))print(x.auto_grad) # tensor([-0.0121,  0.0055, -0.0756, -0.0747,  0.0134,  0.0867, -0.0546,  0.1121, -0.0934, -0.1046], grad_fn=)
mlp = MLP()mlp.load_state_dict(mlp_state_dict)y = mlp(x)z = torch.sum(y)z.backward()print(x.grad) # tensor([-0.0121,  0.0055, -0.0756, -0.0747,  0.0134,  0.0867, -0.0546,  0.1121, -0.0934, -0.1046])

PyTorch 使用的是动态图，而它的计算图在每次前向传播时都是从头开始构建，所以它能够让我们根据需求使用 python 控制语句（如 for、if 等）创建计算图。下面提供一个例子：

def f(x):    result = 1    for ii in x:        if ii.item()>0: result=ii*result    return result
x = torch.tensor([0.3071,  1.1043,  1.3605, -0.3471], requires_grad=True)y = f(x)    # y = x[0]*x[1]*x[2]y.backward()print(x.grad)   # tensor([1.5023, 0.4178, 0.3391, 0.0000])
x = torch.tensor([ 1.2817,  1.7840, -1.7033,  0.1302], requires_grad=True)y = f(x)    # y = x[0]*x[1]*x[3]y.backward()print(x.grad)   # tensor([0.2323, 0.1669, 0.0000, 2.2866])

此前的例子使用的是 Tensor.backward() 接口（内部调用 autograd.backward），接下来我们来介绍 autograd 提供的 jacobian() 和 hessian() 接口，并直接利用其进行自动微分。这两个函数的输入为运算函数（接受输入 tensor，返回输出 tensor）和输入 tensor，返回 jacobian 和 hessian 矩阵。对于jacobian 接口，输入输出均可以为 n 维张量；对于 hessian 接口，输出必需为一标量。

jacobian 返回的张量 shape 为 output_dim x input_dim（若函数输出为标量，则 output_dim 可省略），hessian 返回的张量为 input_dim x input_dim。除此之外，这两个自动微分接口同时支持运算函数接收和输出多个 tensor。

from torch.autograd.functional import jacobian, hessianfrom torch.nn import Linear, AvgPool2d
fc = Linear(4, 2)pool = AvgPool2d(kernel_size=2)
def scalar_func(x):    y = x ** 2    z = torch.sum(y)    return z
def vector_func(x):    y = fc(x)    return y
def mat_func(x):    x = x.reshape((1, 1,) + x.shape)    x = pool(x)    x = x.reshape(x.shape[2:])    return x ** 2
vector_input = torch.randn(4, requires_grad=True)mat_input = torch.randn((4, 4), requires_grad=True)
j = jacobian(scalar_func, vector_input)assert j.shape == (4, )assert torch.all(jacobian(scalar_func, vector_input) == 2 * vector_input)h = hessian(scalar_func, vector_input)assert h.shape == (4, 4)assert torch.all(hessian(scalar_func, vector_input) == 2 * torch.eye(4))j = jacobian(vector_func, vector_input)assert j.shape == (2, 4)assert torch.all(j == fc.weight)j = jacobian(mat_func, mat_input)assert j.shape == (2, 2, 4, 4)

在前面的例子中，我们已经介绍了 autograd.backward() 为节约空间，仅会保存叶节点的梯度。若我们想得知输出关于某一中间结果的梯度，我们则可以选择使用 autograd.grad() 接口，或是使用 hook 机制：

A = torch.tensor(2., requires_grad=True)B = torch.tensor(.5, requires_grad=True)C = A * BD = C.exp()torch.autograd.grad(D, (C, A))  # (tensor(2.7183), tensor(1.3591)), 返回的梯度为tuple类型, grad接口支持对多个变量计算梯度
def variable_hook(grad):                        # hook注册在Tensor上，输入为反传至这一tensor的梯度    print('the gradient of C is：', grad)
A = torch.tensor(2., requires_grad=True)B = torch.tensor(.5, requires_grad=True)C = A * Bhook_handle = C.register_hook(variable_hook)    # 在中间变量C上注册hookD = C.exp()                 D.backward()                                    # 反传时打印：the gradient of C is：tensor(2.7183)hook_handle.remove()                            # 如不再需要，可remove掉这一hook

数值梯度检查

torch.autograd.gradcheck

在编写好自己的 autograd function 后，可以利用 gradcheck 中提供的 gradcheck 和 gradgradcheck 接口，对数值算得的梯度和求导算得的梯度进行比较，以检查 backward 是否编写正确。以函数 $f(x) = y, x \in R^n, y \in R$ 为例，数值法求得 $x_0$ 点的梯度为：

$\frac{\partial y}{\partial x_i} = \frac{f(x_0 + \epsilon e_i) - f(x_0 - \epsilon e_i)}{2 \epsilon}$

在下面的例子中，我们自己实现了 Sigmoid 函数，并利用 gradcheck 来检查 backward 的编写是否正确。

class Sigmoid(Function):                                                    @staticmethod    def forward(ctx, x):         output = 1 / (1 + torch.exp(-x))        ctx.save_for_backward(output)        return output
    @staticmethod    def backward(ctx, grad_output):         output,  = ctx.saved_tensors        grad_x = output * (1 - output) * grad_output        return grad_x
test_input = torch.randn(4, requires_grad=True)     # tensor([-0.4646, -0.4403,  1.2525, -0.5953], requires_grad=True)torch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-3)    # passtorch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-3)    # passtorch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-4)    # failtorch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-4)    # fail

我们发现：eps 为 1e-3 时，我们编写的 Sigmoid 和 torch 自带的 builtin Sigmoid 都可以通过梯度检查；但 eps 下降至 1e-4 时，两者反而都无法通过。而一般我们的直觉是，计算数值梯度时， eps 越小，求得的值应该更接近于真实的梯度。那么为什么这里会出现反常的现象呢？

其实这是由于机器精度带来的误差所致：test_input 的类型为 torch.float32，因此在 eps 过小的情况下，产生了较大的精度误差（计算数值梯度时，eps 作为被除数），因而与真实精度间产生了较大的 gap。当我们将 test_input 换为 float64 的 tensor 后，将不再出现这一现象。

这个现象也为我们提了个醒：在编写 backward 时，要考虑数值计算的一些性质，尽可能保留更精确的结果。

test_input = torch.randn(4, requires_grad=True, dtype=torch.float64)    # tensor([-0.4646, -0.4403,  1.2525, -0.5953], dtype=torch.float64, requires_grad=True)torch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-4)    # passtorch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-4)    # pass
torch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-6)    # passtorch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-6)    # pass

在自动求导时检测错误产生路径

torch.autograd.anomaly_mode

在这个过程中，还有一个问题我们不容忽视：如何在自动求导时检测错误产生路径呢？一般情况下，我们可借助 with autograd.detect_anomaly() 或是 torch.autograd.set_detect_anomaly(True) 来启用：

>>> import torch>>> from torch import autograd>>>>>> class MyFunc(autograd.Function):......     @staticmethod...     def forward(ctx, inp):...         return inp.clone()......     @staticmethod...     def backward(ctx, gO):...         # Error during the backward pass...         raise RuntimeError("Some error in backward")...         return gO.clone()>>>>>> def run_fn(a):...     out = MyFunc.apply(a)...     return out.sum()>>>>>> inp = torch.rand(10, 10, requires_grad=True)>>> out = run_fn(inp)>>> out.backward()    Traceback (most recent call last):      Some Error Log    RuntimeError: Some error in backward>>> with autograd.detect_anomaly():...     inp = torch.rand(10, 10, requires_grad=True)...     out = run_fn(inp)...     out.backward()    Traceback of forward call that caused the error:            # 检测到错误发生的Trace      File "tmp.py", line 53, in         out = run_fn(inp)      File "tmp.py", line 44, in run_fn        out = MyFunc.apply(a)    Traceback (most recent call last):      Some Error Log    RuntimeError: Some error in backward

设置是否需要梯度

torch.autograd.grad_mode

我们在 inference 的过程中，不希望 autograd 对 tensor 求导，因为求导需要缓存许多中间结构，增加额外的内存/显存开销。在 inference 时，关闭自动求导可实现一定程度的速度提升，并节省大量内存及显存（被节省的不仅限于原先用于梯度存储的部分。torch 中我们可以利用 grad_mode 中的 troch.no_grad() 来关闭自动求导：

from torchvision.models import resnet50import torch
net = resnet50().cuda(0)num = 128inp = torch.ones([num, 3, 224, 224]).cuda(0)net(inp)                                        # 若不开torch.no_grad()，batch_size为128时就会OOM (在1080 Ti上)
net = resnet50().cuda(1)num = 512inp = torch.ones([num, 3, 224, 224]).cuda(1)    with torch.no_grad():                           # 打开torch.no_grad()后，batch_size为512时依然能跑inference (节约超过4倍显存)    net(inp)

model.eval() 与 torch.no_grad()

这两项在 inference 的过程中需要都打开：model.eval() 令 model 中的 BatchNorm, Dropout 等 module 采用 eval mode，保证 inference 结果的正确性，但不起到节省显存的作用；torch.no_grad() 声明不计算梯度，可节省大量内存和显存。

提供 function 级别的统计信息

torch.autograd.profiler

import torchfrom torchvision.models import resnet18
x = torch.randn((1, 3, 224, 224), requires_grad=True)model = resnet18()with torch.autograd.profiler.profile() as prof:    for _ in range(100):        y = model(x)        y = torch.sum(y)        y.backward()# NOTE: some columns were removed for brevityprint(prof.key_averages().table(sort_by="self_cpu_time_total"))

输出为包含 CPU 时间及占比，调用次数等信息。这里需要注意：由于一个 kernel 可能还会调用其他 kernel，因此 Self CPU 指 CPU 本身所耗时间（不含其他 kernel 被调用所耗时间）。

下表展示了使用一块联想小新 i5 轻薄本所耗的时间，我们可以看到：耗时达到了惊人的 51s～

-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------              aten::mkldnn_convolution_backward_weights        26.85%       13.936s        26.99%       14.012s       7.006ms          2000                aten::mkldnn_convolution_backward_input        24.20%       12.564s        24.28%       12.605s       6.302ms          2000                               aten::mkldnn_convolution        22.66%       11.764s        22.75%       11.810s       5.905ms          2000                                             aten::add_         4.54%        2.358s         4.54%        2.358s     335.031us          7037                                aten::native_batch_norm         4.20%        2.178s         6.35%        3.294s       1.647ms          2000                               aten::threshold_backward         3.85%        1.999s         3.85%        1.999s       1.176ms          1700                       aten::native_batch_norm_backward         2.40%        1.247s         2.60%        1.348s     673.925us          2000                 struct torch::autograd::AccumulateGrad         1.75%     906.138ms         5.90%        3.062s     485.973us          6300                                              aten::sum         1.56%     808.045ms         1.61%     836.033ms     363.493us          2300                          aten::max_pool2d_with_indices         1.26%     655.751ms         1.26%     655.751ms       6.558ms           100                                              aten::add         0.78%     407.001ms         0.78%     407.001ms     145.357us          2800                                        aten::clamp_min         0.71%     367.014ms         0.71%     367.014ms     215.891us          1700autograd::engine::evaluate_function: NativeBatchNorm...         0.55%     287.113ms         3.29%        1.707s     853.746us          2000                                            aten::empty         0.49%     255.215ms         0.49%     255.225ms       8.451us         30200autograd::engine::evaluate_function: MkldnnConvoluti...         0.47%     246.019ms        52.59%       27.302s      13.651ms          2000                                               aten::mm         0.45%     232.766ms         0.46%     240.850ms       1.204ms           200     autograd::engine::evaluate_function: ReluBackward0         0.40%     206.792ms         4.35%        2.256s       1.327ms          1700                                            aten::fill_         0.30%     154.547ms         0.30%     154.547ms      61.819us          2500                 aten::max_pool2d_with_indices_backward         0.24%     122.017ms         0.50%     260.168ms       2.602ms           100autograd::engine::evaluate_function: struct torch::a...         0.20%     106.389ms         6.10%        3.168s     502.860us          6300autograd::engine::evaluate_function: MaxPool2DWithIn...         0.16%      84.665ms         0.76%     395.121ms       3.951ms           100                                       aten::empty_like         0.15%      75.751ms         0.30%     154.099ms      18.566us          8300                               NativeBatchNormBackward0         0.14%      72.529ms         2.74%        1.420s     710.189us          2000                                             aten::div_         0.14%      70.880ms         0.33%     172.975ms      82.369us          2100                                            aten::addmm         0.13%      68.586ms         0.14%      71.177ms     711.770us           100                      aten::mkldnn_convolution_backward         0.12%      61.287ms        51.39%       26.678s      13.339ms          2000                                     aten::_convolution         0.11%      57.955ms        22.86%       11.868s       5.934ms          2000                                             aten::mean         0.11%      57.625ms         2.03%        1.054s     501.744us          2100                          MaxPool2DWithIndicesBackward0         0.10%      50.288ms         0.60%     310.456ms       3.105ms           100                                          ReluBackward0         0.10%      49.924ms         3.95%        2.049s       1.205ms          1700                                            aten::relu_         0.09%      47.010ms         0.83%     430.029ms     252.958us          1700                           aten::_batch_norm_impl_index         0.09%      45.831ms         6.45%        3.349s       1.674ms          2000                                         aten::_to_copy         0.07%      37.126ms         0.15%      77.273ms      35.124us          2200                                      aten::as_strided_         0.06%      32.698ms         0.06%      32.698ms       5.450us          6000                             MkldnnConvolutionBackward0         0.06%      29.531ms        51.45%       26.707s      13.354ms          2000                                               aten::to         0.06%      29.155ms         0.21%     106.428ms      46.273us          2300                                           aten::conv2d         0.05%      28.411ms        22.95%       11.914s       5.957ms          2000                                            aten::copy_         0.05%      28.014ms         0.05%      28.014ms      11.206us          2500                                       aten::batch_norm         0.05%      24.090ms         6.50%        3.373s       1.686ms          2000                                              aten::div         0.04%      21.886ms         0.05%      26.064ms     260.640us           100                                      aten::convolution         0.03%      18.080ms        22.90%       11.886s       5.943ms          2000                                    aten::empty_strided         0.03%      16.717ms         0.03%      16.717ms       7.268us          2300                                       aten::clamp_min_         0.03%      16.005ms         0.74%     383.019ms     225.305us          1700                                       aten::as_strided         0.03%      13.844ms         0.03%      13.844ms       4.774us          2900      autograd::engine::evaluate_function: AddBackward0         0.02%      12.168ms         0.03%      14.762ms      18.453us           800                                                aten::t         0.01%       7.523ms         0.03%      13.994ms      27.988us           500                                   aten::_reshape_alias         0.01%       6.235ms         0.01%       6.235ms      31.175us           200                                         AddmmBackward0         0.01%       5.238ms         0.49%     253.384ms       2.534ms           100    autograd::engine::evaluate_function: AddmmBackward0         0.01%       4.729ms         0.51%     265.184ms       2.652ms           100                                           aten::expand         0.01%       4.448ms         0.01%       5.570ms      18.567us           300                                        aten::transpose         0.01%       4.039ms         0.01%       6.471ms      12.942us           500                                           AddBackward0         0.00%       2.594ms         0.00%       2.594ms       3.243us           800                                       aten::max_pool2d         0.00%       2.513ms         1.27%     658.264ms       6.583ms           100                                          MeanBackward1         0.00%       2.391ms         0.06%      30.648ms     306.480us           100                                            aten::clone         0.00%       2.388ms         0.01%       7.760ms      38.800us           200                                          aten::flatten         0.00%       2.179ms         0.01%       6.422ms      64.220us           100                                           aten::linear         0.00%       2.104ms         0.15%      77.744ms     777.440us           100                                             aten::view         0.00%       1.449ms         0.00%       1.449ms      14.490us           100                                            aten::zero_         0.00%       1.418ms         0.27%     138.151ms       1.382ms           100                                  ReshapeAliasBackward0         0.00%       1.323ms         0.01%       4.498ms      44.980us           100     autograd::engine::evaluate_function: MeanBackward1         0.00%       1.295ms         0.06%      31.943ms     319.430us           100                              aten::adaptive_avg_pool2d         0.00%       1.283ms         0.05%      25.883ms     258.830us           100                                          aten::reshape         0.00%       1.183ms         0.01%       3.175ms      31.750us           100autograd::engine::evaluate_function: ReshapeAliasBac...         0.00%       1.177ms         0.01%       5.675ms      56.750us           100                                           SumBackward0         0.00%       1.148ms         0.01%       3.796ms      37.960us           100        autograd::engine::evaluate_function: TBackward0         0.00%       1.073ms         0.01%       3.808ms      38.080us           100      autograd::engine::evaluate_function: SumBackward0         0.00%       1.008ms         0.01%       4.804ms      48.040us           100                                        aten::ones_like         0.00%     946.000us         0.01%       2.688ms      26.880us           100                                                 detach         0.00%     512.000us         0.00%     512.000us       8.127us            63                                     aten::resolve_conj         0.00%     505.000us         0.00%     505.000us       0.842us           600                                             TBackward0         0.00%     500.000us         0.01%       2.735ms      27.350us           100                                           aten::detach         0.00%     240.000us         0.00%     752.000us      11.937us            63-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------Self CPU time total: 51.911s

本次的梯度计算模块分享就到此结束啦！恭喜看到这里的你又 get 了一个新技能～在之后的系列分享中，我们会继续给大家带来有关 PyTorch 优化算法接口的相关知识。点个关注不迷路，我们下期见哦！