pytorch统计模型计算量和参数量

利用thop库中的profile函数,可以实现统计pytorch模型的计算量和参数量。

import torch
from thop import profile

net = Net()     #net为pytorch的网络模型
input_ = torch.rand((1,2,512,512))  #随机初始化一个向量,用于网络输入
flops, params = profile(net, inputs = (input_,)) #得到计算量和参数量

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注