用thop去获得模型参数量和计算量（pytorch）

Thop github地址

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