在需要计算代码前使用profiler进行声明,即可输出运行开销
from torch.autograd import Variable
import torch
x = Variable(torch.randn(1, 1), requires_grad=True)
with torch.autograd.profiler.profile(use_cuda=True) as prof:
y = x ** 2
y.backward()
# NOTE: some columns were removed for brevity
print(prof)