pytorch统计模型计算强度

计算强度 = 运算数 / 访存数

运算数有很多库可以算，例如thop的profile

from thop import profile
model = torchvision.models.resnet18(weights=None)
flops, params = profile(model, inputs=(torch.randn(1, 3, 224, 224), ))
print("flops: {:.2f}Gflops".format(flops/1000/1000/1000))

访存数目前只找到了torchstat

from torchstat import stat
model = torchvision.models.resnet18(weights=None)
stat(model, (3, 224, 224))

torchstat使用问题

问题1 vit模型套用会出错

结合报错，发现是vit中存在(1,a,b)这样输入的线性层。但是torchstat中是会报错的。
解决办法：
找到相应的库位置，对compute_flops.py compute_madd.py compute_memory.py三个文件中的进行修改。
compute_Linear_flops compute_Linear_madd compute_Linear_memory三个函数中的len(inp.size()) == 2 and len(out.size()) == 2
都修改为

 assert (len(inp.size()) == 2 and len(out.size()) == 2) or (len(inp.size()) == 3 and inp.size()[0] == 1 and len(out.size()) == 3 and out.size()[0] == 1)
 if len(inp.size()) > 2:
    inp = inp[0]
 if len(out.size()) > 2:
    out = out[0]

问题2 产生报告时，MemRead与MemWrite没进行求和

找到相应的库位置，对reporter.py的61行进行修改。
将mread, mwrite替换为total_mread, total_mwrite

问题3 memory过大导致溢出变负数

一般tensor是不会溢出的，经过检查最后发现是因为numpy的数据转换出现了点问题。
将model_hook.py中的Memory = np.array(Memory, dtype=np.int32) * itemsize
替换为Memory = np.array(Memory, dtype=np.int64) * itemsize
最后总函数

dic = {'vgg16':torchvision.models.vgg16,
        'resnet18':torchvision.models.resnet18,
        'resnet34':torchvision.models.resnet34,
        'resnet50':torchvision.models.resnet50,
        'resnet101':torchvision.models.resnet101,
        'resnet152':torchvision.models.resnet152,
        'vit_b_16':torchvision.models.vit_b_16,
        'vit_b_32':torchvision.models.vit_b_32,
        'vit_l_16':torchvision.models.vit_l_16,
        'vit_l_32':torchvision.models.vit_l_32,
        'vit_h_14':torchvision.models.vit_h_14,
        }
models_vgg = ['vgg16','resnet18','resnet34','resnet50','resnet101','resnet152']
models_vit = ['vit_b_16','vit_b_32','vit_l_16','vit_l_32','vit_h_14']
models = models_vgg + models_vit
for model_name in models_test:
  model = dic[model_name](weights=None)
  stat(model, (3, 224, 224))
  for i in range(5):
    print('')

vit统计问题

统计vit模型时，torchstat与thop都不会统计多头注意力层，所以flops与参数量统计都有问题。需要进行修正。

以vit_b_16为例所向手推（batch_size为1，不保证正确）：
q:(1,197,768)经过(768,768)的线性层得到(1,197,768)，参数量：768*768，flops:197*768*768*2
k:(1,197,768)经过(768,768)的线性层得到(1,197,768)，参数量：768*768，flops:197*768*768*2
v:(1,197,768)经过(768,768)的线性层得到(1,197,768)，参数量：768*768，flops:197*768*768*2
其中，197=(224/16)*(224/16)+1
q,k,v reshape为(12,197,64)，12为head数。
对每一个头来说，(197,64)*(64,197)*(197,64)得到(197,64)，所有头这拼接完为(12,197,64)，reshape成(197,768)
flops:(197*64*197*2+197*197*64*2)*12，即(197*768*197*2+197*197*768*2)。
可以发现与head没有关系。
最后再经过一个(768,768)的线性层，参数量：768*768，flops:197*768*768*2
所以最后要在统计结果上参数量加768*768*4*12（12为layer数量）
flops加197*768*(4*768+2*197)*2*12（12为layer数量）