pytorch操作小技巧----模型性能观察

在我们创造一个新模型的时候，我们需要对一个模型进行全面的观察，与其他模型进行比较。

下面就介绍几种观察模型性能的方法

1查看模型的计算量和参数
使用thop这个参数库

from thop import profile
from thop import clever_format
def get_pf(model,input):
    flops, params = profile(model, inputs=(input,))
    flops, params = clever_format([flops, params], "%.3f")
    return flops, params

调用这个函数将会返回flops以及参数量。

flops通常代表的是浮点运算数，衡量的是算法或者模型的复杂度，计算量小代表计算过程中需要的计算空间较小

参数量就是指模型各个部分的参数数量之和，参数小指模型的参数少，比较轻量级

但是两者并不是并行的，意思就是计算量小不代表参数量小，参数量小也不代表计算量就少，所以我们要两者结合在一起判断。

具体调用函数的例子可以如下

model=resnet50().cuda()
inputs=torch.rand(4,3,256,128).cuda()
get_pf(model,inputs)

查看模型各部分的计算耗时和内存占用
这个功能已经在torch1.9可以方便的调用

from torch.profiler import profile, record_function, ProfilerActivity
model=resnet50().cuda()
inputs=torch.rand(4,3,256,128).cuda()

with profile(activities=[
        ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memory=True,record_shapes=True) as prof:
    with record_function("model_inference"):
        model(inputs)

通过上面的代码将会计算在输入为inputs下，model各个模块在cpu以及cuda上的耗时和内存占用

我们可以根据耗时或者内存进行排序，然后查看哪些模块耗时最大或者内存占用最多，然后进一步针对性的修改模型

根据耗时排序

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

根据内存占用排序

print(prof.key_averages().table(sort_by="self_cuda_mem", row_limit=10))

原文链接：https://zhuanlan.zhihu.com/p/393840388