1. 安装
pip install thop
2. 使用（获得计算量（flops）和参数量（params））
2.1计算tochvision自带的模型大小
其中输入的x，只影响计算量，不影响参数量
model = torchvision.models.segmentation.fcn_resnet50() x = torch.randn(1,3,224,224) flops, params = thop.profile(model,inputs=(x,))
2.2计算自定义的模型的大小
传入的模型一定要是class类型的
我们自己定义的一个模型（Unet）
class U_Net(nn.Module): """ UNet - Basic Implementation Paper : https://arxiv.org/abs/1505.04597 """ def __init__(self, in_ch=3, out_ch=1): super(U_Net, self).__init__() n1 = 64 filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16] self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2) self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2) self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2) self.Conv1 = conv_block(in_ch, filters[0]) self.Conv2 = conv_block(filters[0], filters[1]) self.Conv3 = conv_block(filters[1], filters[2]) self.Conv4 = conv_block(filters[2], filters[3]) self.Conv5 = conv_block(filters[3], filters[4]) self.Up5 = up_conv(filters[4], filters[3]) self.Up_conv5 = conv_block(filters[4], filters[3]) self.Up4 = up_conv(filters[3], filters[2]) self.Up_conv4 = conv_block(filters[3], filters[2]) self.Up3 = up_conv(filters[2], filters[1]) self.Up_conv3 = conv_block(filters[2], filters[1]) self.Up2 = up_conv(filters[1], filters[0]) self.Up_conv2 = conv_block(filters[1], filters[0]) self.Conv = nn.Conv2d(filters[0], out_ch, kernel_size=1, stride=1, padding=0) # self.active = torch.nn.Sigmoid() def forward(self, x): e1 = self.Conv1(x) e2 = self.Maxpool1(e1) e2 = self.Conv2(e2) e3 = self.Maxpool2(e2) e3 = self.Conv3(e3) e4 = self.Maxpool3(e3) e4 = self.Conv4(e4) e5 = self.Maxpool4(e4) e5 = self.Conv5(e5) d5 = self.Up5(e5) d5 = torch.cat((e4, d5), dim=1) d5 = self.Up_conv5(d5) d4 = self.Up4(d5) d4 = torch.cat((e3, d4), dim=1) d4 = self.Up_conv4(d4) d3 = self.Up3(d4) d3 = torch.cat((e2, d3), dim=1) d3 = self.Up_conv3(d3) d2 = self.Up2(d3) d2 = torch.cat((e1, d2), dim=1) d2 = self.Up_conv2(d2) out = self.Conv(d2) # d1 = self.active(out) return out
计算模型大小
model = U_Net() x = torch.randn(1,3,224,224) flops, params = thop.profile(model,inputs=(x,))
最后的flops和params就是计算量和参数量了。
下面是利用thop计算了，一些图像分割模型的计算量(左)和参数量(右)。
26535694416.0 fcn_res50 32957013.0 41425522768.0 fcn_res101 51949141.0 50123064320.0 U_net 34527041.0 50966421744.0 Att_Unet 34878573.0 116999443456.0 R2_Unet 39091393.0 117842800880.0 R2AttU_Net 39442925.0 106110247936.0 NestedUnet 36629633.0
