PyTorch 原始碼解讀之 torch.autograd
阿新 • • 發佈：2020-12-21
AI編輯：我是小將
本文作者：OpenMMLab
https://zhuanlan.zhihu.com/p/321449610
本文已由原作者授權
前言
本篇筆記以介紹 pytorch 中的 autograd 模組功能為主，主要涉及 torch/autograd 下程式碼，不涉及底層的 C++ 實現。本文涉及的原始碼以 PyTorch 1.7 為準。
-
torch.autograd.function （函式的反向傳播）
-
torch.autograd.functional （計算圖的反向傳播）
-
torch.autograd.gradcheck （數值梯度檢查）
-
torch.autograd.anomaly_mode （在自動求導時檢測錯誤產生路徑）
-
torch.autograd.grad_mode （設定是否需要梯度）
-
model.eval() 與 torch.no_grad()
-
torch.autograd.profiler （提供 function 級別的統計資訊）
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 function
print(ret) # tensor([2.7183], grad_fn=<ExpBackward>)
ret.backward() # 反傳梯度
print(x.grad) # tensor([2.7183])Exp 函式的前向很簡單，直接呼叫 tensor 的成員方法exp即可。反向時，我們知道, 因此我們直接使用乘以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.1
ret = ret ** 2
print(ret) # tensor([4.], grad_fn=<PowBackward0>)
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 * B
D = C.exp()
F = D + E
print(F) # tensor(3.7183, grad_fn=<AddBackward0>) 列印計算結果，可以看到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]]) # [<AddBackward0 object at 0x7f972de8c7b8>, <ExpBackward object at 0x7f972de8c278>, <MulBackward0 object at 0x7f972de8c2b0>, None] 每個變數的grad_fn指向產生其運算元的backward function，葉節點的grad_fn為空
print(F.grad_fn.next_functions) # ((<ExpBackward object at 0x7f972de8c390>, 0), (<AccumulateGrad object at 0x7f972de8c5f8>, 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 ** 2
D = B ** 2
E = C * D
F = 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=<MulBackward0>), tensor([40.], grad_fn=<MulBackward0>))下面，我們編寫一個簡單的函式，在這個計算圖上進行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 ** 2
D = B ** 2
E = C * D
F = D + E
autograd(F.grad_fn, torch.tensor(1))
print(A.auto_grad, B.auto_grad) # tensor(24., grad_fn=<UnbindBackward>) tensor(40., grad_fn=<AddBackward0>)
# 這一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()
# 自定義autograd
mlp = 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=<AddBackward0>)
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, hessian
from 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 * B
D = 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 * B
hook_handle = C.register_hook(variable_hook) # 在中間變數C上註冊hook
D = C.exp()
D.backward() # 反傳時列印：the gradient of C is：tensor(2.7183)
hook_handle.remove() # 如不再需要，可remove掉這一hooktorch.autograd.gradcheck（數值梯度檢查）
在編寫好自己的 autograd function 後，可以利用gradcheck中提供的gradcheck和gradgradcheck介面，對數值算得的梯度和求導算得的梯度進行比較，以檢查backward是否編寫正確。以函式為例，數值法求得點的梯度為：。在下面的例子中，我們自己實現了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) # pass
torch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-3) # pass
torch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-4) # fail
torch.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) # pass
torch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-4) # pass
torch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-6) # pass
torch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-6) # passtorch.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 <module>
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 backwardtorch.autograd.grad_mode（設定是否需要梯度）
我們在 inference 的過程中，不希望 autograd 對 tensor 求導，因為求導需要快取許多中間結構，增加額外的記憶體/視訊記憶體開銷。在 inference 時，關閉自動求導可實現一定程度的速度提升，並節省大量記憶體及視訊記憶體（被節省的不僅限於原先用於梯度儲存的部分）。我們可以利用grad_mode中的troch.no_grad()來關閉自動求導：
from torchvision.models import resnet50
import torch
net = resnet50().cuda(0)
num = 128
inp = torch.ones([num, 3, 224, 224]).cuda(0)
net(inp) # 若不開torch.no_grad()，batch_size為128時就會OOM (在1080 Ti上)
net = resnet50().cuda(1)
num = 512
inp = 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()宣告不計算梯度，節省大量記憶體和視訊記憶體。
torch.autograd.profiler（提供function級別的統計資訊）
import torch
from 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 brevity
print(prof.key_averages().table(sort_by="self_cpu_time_total"))輸出為包含 CPU 時間及佔比，呼叫次數等資訊（由於一個 kernel 可能還會呼叫其他 kernel，因此 Self CPU 指他本身所耗時間（不含其他 kernel 被呼叫所耗時間））：
--------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg # of Calls
--------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
aten::mkldnn_convolution_backward_input 18.69% 1.722s 18.88% 1.740s 870.001us 2000
aten::mkldnn_convolution 17.07% 1.573s 17.28% 1.593s 796.539us 2000
aten::mkldnn_convolution_backward_weights 16.96% 1.563s 17.21% 1.586s 792.996us 2000
aten::native_batch_norm 9.51% 876.994ms 15.06% 1.388s 694.049us 2000
aten::max_pool2d_with_indices 9.47% 872.695ms 9.48% 873.802ms 8.738ms 100
aten::select 7.00% 645.298ms 10.06% 926.831ms 7.356us 126000
aten::native_batch_norm_backward 6.67% 614.718ms 12.16% 1.121s 560.466us 2000
aten::as_strided 3.07% 282.885ms 3.07% 282.885ms 2.229us 126900
aten::add_ 2.85% 262.832ms 2.85% 262.832ms 37.350us 7037
aten::empty 1.23% 113.274ms 1.23% 113.274ms 4.089us 27700
aten::threshold_backward 1.10% 101.094ms 1.17% 107.383ms 63.166us 1700
aten::add 0.88% 81.476ms 0.99% 91.350ms 32.625us 2800
aten::max_pool2d_with_indices_backward 0.86% 79.174ms 1.02% 93.706ms 937.064us 100
aten::threshold_ 0.56% 51.678ms 0.56% 51.678ms 30.399us 1700
torch::autograd::AccumulateGrad 0.40% 36.909ms 2.81% 258.754ms 41.072us 6300
aten::empty_like 0.35% 32.532ms 0.63% 57.630ms 6.861us 8400
NativeBatchNormBackward 0.32% 29.572ms 12.48% 1.151s 575.252us 2000
aten::_convolution 0.31% 28.182ms 17.63% 1.625s 812.258us 2000
aten::mm 0.27% 24.983ms 0.32% 29.522ms 147.611us 200
aten::stride 0.27% 24.665ms 0.27% 24.665ms 0.583us 42300
aten::mkldnn_convolution_backward 0.22% 20.025ms 36.33% 3.348s 1.674ms 2000
MkldnnConvolutionBackward 0.21% 19.112ms 36.53% 3.367s 1.684ms 2000
aten::relu_ 0.20% 18.611ms 0.76% 70.289ms 41.346us 1700
aten::_batch_norm_impl_index 0.16% 14.298ms 15.32% 1.413s 706.254us 2000
aten::addmm 0.14% 12.684ms 0.15% 14.138ms 141.377us 100
aten::fill_ 0.14% 12.672ms 0.14% 12.672ms 21.120us 600
ReluBackward1 0.13% 11.845ms 1.29% 119.228ms 70.134us 1700
aten::as_strided_ 0.13% 11.674ms 0.13% 11.674ms 1.946us 6000
aten::div 0.11% 10.246ms 0.13% 12.288ms 122.876us 100
aten::batch_norm 0.10% 8.894ms 15.42% 1.421s 710.700us 2000
aten::convolution 0.08% 7.478ms 17.71% 1.632s 815.997us 2000
aten::sum 0.08% 7.066ms 0.10% 9.424ms 31.415us 300
aten::conv2d 0.07% 6.851ms 17.78% 1.639s 819.423us 2000
aten::contiguous 0.06% 5.597ms 0.06% 5.597ms 0.903us 6200
aten::copy_ 0.04% 3.759ms 0.04% 3.980ms 7.959us 500
aten::t 0.04% 3.526ms 0.06% 5.561ms 11.122us 500
aten::view 0.03% 2.611ms 0.03% 2.611ms 8.702us 300
aten::div_ 0.02% 1.973ms 0.04% 4.051ms 40.512us 100
aten::expand 0.02% 1.720ms 0.02% 2.225ms 7.415us 300
AddmmBackward 0.02% 1.601ms 0.37% 34.141ms 341.414us 100
aten::to 0.02% 1.596ms 0.04% 3.871ms 12.902us 300
aten::mean 0.02% 1.485ms 0.10% 9.204ms 92.035us 100
AddBackward0 0.01% 1.381ms 0.01% 1.381ms 1.726us 800
aten::transpose 0.01% 1.297ms 0.02% 2.035ms 4.071us 500
aten::empty_strided 0.01% 1.163ms 0.01% 1.163ms 3.877us 300
MaxPool2DWithIndicesBackward 0.01% 1.095ms 1.03% 94.802ms 948.018us 100
MeanBackward1 0.01% 974.822us 0.16% 14.393ms 143.931us 100
aten::resize_ 0.01% 911.689us 0.01% 911.689us 3.039us 300
aten::zeros_like 0.01% 884.496us 0.11% 10.384ms 103.843us 100
aten::clone 0.01% 798.993us 0.04% 3.687ms 18.435us 200
aten::reshape 0.01% 763.804us 0.03% 2.604ms 13.021us 200
aten::zero_ 0.01% 689.598us 0.13% 11.919ms 59.595us 200
aten::resize_as_ 0.01% 562.349us 0.01% 776.967us 7.770us 100
aten::max_pool2d 0.01% 492.109us 9.49% 874.295ms 8.743ms 100
aten::adaptive_avg_pool2d 0.01% 469.736us 0.10% 9.673ms 96.733us 100
aten::ones_like 0.00% 460.352us 0.01% 1.377ms 13.766us 100
SumBackward0 0.00% 399.188us 0.01% 1.206ms 12.057us 100
aten::flatten 0.00% 397.053us 0.02% 1.917ms 19.165us 100
ViewBackward 0.00% 351.824us 0.02% 1.436ms 14.365us 100
TBackward 0.00% 308.947us 0.01% 1.315ms 13.150us 100
detach 0.00% 127.329us 0.00% 127.329us 2.021us 63
torch::autograd::GraphRoot 0.00% 114.731us 0.00% 114.731us 1.147us 100
aten::detach 0.00% 106.170us 0.00% 233.499us 3.706us 63
--------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 9.217sReference
[1]Automatic differentiation package - torch.autograd — PyTorch 1.7.0 documentation
[2]Autograd
推薦閱讀
MMDetection新版本V2.7釋出，支援DETR，還有YOLOV4在路上!
無需tricks，知識蒸餾提升ResNet50在ImageNet上準確度至80%+
不妨試試MoCo，來替換ImageNet上pretrain模型！