NORTON/complexity.py at main · vantienpham/NORTON · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import argparse
import torch
from models.cifar10 import *
from models.imagenet import *
from detection.model import (
    fasterrcnn_CPresnet50_fpn,
    maskrcnn_CPresnet50_fpn,
    keypointrcnn_CPresnet50_fpn,
)
from utils.common import get_cpr
from ptflops import get_model_complexity_info
from thop import profile


def parse_args():
    parser = argparse.ArgumentParser("Compute model complexity")

    parser.add_argument(
        "--dataset",
        type=str,
        default="cifar10",
        choices=("cifar10", "imagenet", "coco"),
        help="dataset",
    )
    parser.add_argument(
        "--arch",
        type=str,
        default="vgg_16_bn",
        choices=(
            "vgg_16_bn",
            "vgg_19_bn",
            "resnet_56",
            "resnet_110",
            "densenet_40",
            "resnet_50",
            "fasterrcnn_CPresnet50_fpn",
            "maskrcnn_CPresnet50_fpn",
            "keypointrcnn_CPresnet50_fpn",
        ),
        help="architecture",
    )
    parser.add_argument(
        "-r",
        "--rank",
        dest="rank",
        type=int,
        default=0,
        help="use pre-specified rank for all layers",
    )
    parser.add_argument(
        "-cpr",
        "--compress_rate",
        type=str,
        default="[0.]*100",
        help="list of compress rate of each layer",
    )

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    with torch.cuda.device(0):
        compress_rate = get_cpr(args.compress_rate)
        model = eval(args.arch)(compress_rate=compress_rate, rank=args.rank)

        inp_img_size = 32
        if args.dataset == "imagenet":
            inp_img_size = 224
        elif args.dataset == "coco":
            inp_img_size = 800

        macs_ptfl, params_ptfl = get_model_complexity_info(
            model,
            (3, inp_img_size, inp_img_size),
            as_strings=False,
            print_per_layer_stat=False,
            verbose=False,
        )

        input = torch.randn(1, 3, inp_img_size, inp_img_size)
        macs_thop, params_thop = profile(model, inputs=(input,))

        macs = min(macs_ptfl, macs_thop)
        params = min(params_ptfl, params_thop)

        print("{:<30}  {:<8}".format("Computational complexity: ", macs))
        print("{:<30}  {:<8}".format("Number of parameters: ", params))

        ori_model = eval(args.arch)(compress_rate=[0.0] * 100, rank=0)
        ori_macs, ori_params = get_model_complexity_info(
            ori_model,
            (3, inp_img_size, inp_img_size),
            as_strings=False,
            print_per_layer_stat=False,
            verbose=False,
        )
        mac_reduced = (1 - macs / ori_macs) * 100
        param_reduced = (1 - params / ori_params) * 100
        print(f"FLOPs_reduced = {mac_reduced:.2f}")
        print(f"param_reduced = {param_reduced:.2f}")

        # for cpr in list(['[0.05]*7+[0.2]*6', '[0.2]*7+[0.5]*6', '[0.25]*7+[0.75]*6']):
        #     compress_rate = get_cpr(cpr)
        #     print(f'compress_rate {compress_rate}')
        #     for rank in range(1, 9):
        #         model = eval(args.arch)(compress_rate, rank)

        #         macs, params = get_model_complexity_info(model,
        #                                          (3, inp_img_size, inp_img_size),
        #                                          as_strings=False,
        #                                          print_per_layer_stat=False,
        #                                          verbose=False)

        #         mac_reduced = (1 - macs/ori_macs)*100
        #         param_reduced = (1- params/ori_params)*100
        #         print(f'rank = {rank}, FLOPs_reduced = {mac_reduced:.2f}')
        #         print(f'rank = {rank}, param_reduced = {param_reduced:.2f}')