首先导入相关库

from torch.profiler import profile, ProfilerActivity

接下来使用下列语句，'your code'处放置自己的代码

with profile(
     activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
     record_shapes=True
) as prof:
     'your code'
print(prof.key_averages().table(sort_by="cpu_time_total"))

如果分析对象是一个函数，也可以将上述语句放入装饰器中，应用起来更加方便

def profile_decorator(func):
    def wrapper(*args, **kwargs):
        with profile(
                activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
                record_shapes=True
        ) as prof:
            result = func(*args, **kwargs)
        print(prof.key_averages().table(sort_by="cpu_time_total"))
        return result
    return wrapper

然后在目标函数处应用装饰器即可

@profile_decorator
def your_function(self, args):
    your code
    bala
    bala

随后运行程序，就会在控制台获得输出，如下所示

可以看到代码每一环节所消耗的cpu/gpu时间和所占的比例