model_speedup.py 3.43 KB
Newer Older
1
import os
2
import sys
QuanluZhang's avatar
QuanluZhang committed
3
4
5
import argparse
import time
import torch
6

7
from pathlib import Path
J-shang's avatar
J-shang committed
8
sys.path.append(str(Path(__file__).absolute().parents[2] / 'models'))
9
10
11
from cifar10.vgg import VGG
from mnist.lenet import LeNet

liuzhe-lz's avatar
liuzhe-lz committed
12
from nni.compression.pytorch import apply_compression_results, ModelSpeedup
QuanluZhang's avatar
QuanluZhang committed
13
14

torch.manual_seed(0)
15
16
17
use_mask = True
use_speedup = True
compare_results = True
QuanluZhang's avatar
QuanluZhang committed
18

19
20
config = {
    'apoz': {
J-shang's avatar
J-shang committed
21
22
23
        'model_name': 'vgg16',
        'input_shape': [64, 3, 32, 32],
        'masks_file': './experiment_data/mask_vgg16_cifar10_apoz.pth'
24
25
26
27
    },
    'l1filter': {
        'model_name': 'vgg16',
        'input_shape': [64, 3, 32, 32],
28
        'masks_file': './experiment_data/mask_vgg16_cifar10_l1filter.pth'
29
30
    },
    'fpgm': {
31
32
33
        'model_name': 'vgg16',
        'input_shape': [64, 3, 32, 32],
        'masks_file': './experiment_data/mask_vgg16_cifar10_fpgm.pth'
34
35
36
37
    },
    'slim': {
        'model_name': 'vgg19',
        'input_shape': [64, 3, 32, 32],
38
        'masks_file': './experiment_data/mask_vgg19_cifar10_slim.pth'
39
40
    }
}
QuanluZhang's avatar
QuanluZhang committed
41

42
43
def model_inference(config):
    masks_file = config['masks_file']
44
45
    device = torch.device(
        'cuda') if torch.cuda.is_available() else torch.device('cpu')
46

47
    # device = torch.device(config['device'])
48
49
50
51
    if config['model_name'] == 'vgg16':
        model = VGG(depth=16)
    elif config['model_name'] == 'vgg19':
        model = VGG(depth=19)
52
53
54
    elif config['model_name'] == 'lenet':
        model = LeNet()

QuanluZhang's avatar
QuanluZhang committed
55
56
57
    model.to(device)
    model.eval()

58
59
60
    dummy_input = torch.randn(config['input_shape']).to(device)
    use_mask_out = use_speedup_out = None
    # must run use_mask before use_speedup because use_speedup modify the model
QuanluZhang's avatar
QuanluZhang committed
61
    if use_mask:
62
        apply_compression_results(model, masks_file, device)
QuanluZhang's avatar
QuanluZhang committed
63
64
        start = time.time()
        for _ in range(32):
65
66
67
            use_mask_out = model(dummy_input)
        print('elapsed time when use mask: ', time.time() - start)
    if use_speedup:
68
        m_speedup = ModelSpeedup(model, dummy_input, masks_file, device)
QuanluZhang's avatar
QuanluZhang committed
69
70
71
        m_speedup.speedup_model()
        start = time.time()
        for _ in range(32):
72
73
74
75
76
77
78
            use_speedup_out = model(dummy_input)
        print('elapsed time when use speedup: ', time.time() - start)
    if compare_results:
        if torch.allclose(use_mask_out, use_speedup_out, atol=1e-07):
            print('the outputs from use_mask and use_speedup are the same')
        else:
            raise RuntimeError('the outputs from use_mask and use_speedup are different')
QuanluZhang's avatar
QuanluZhang committed
79
80
81
82
83
84

if __name__ == '__main__':
    parser = argparse.ArgumentParser("speedup")
    parser.add_argument("--example_name", type=str, default="slim", help="the name of pruning example")
    parser.add_argument("--masks_file", type=str, default=None, help="the path of the masks file")
    args = parser.parse_args()
85
86
87
88
89
90
91
92
93

    if args.example_name != 'all':
        if args.masks_file is not None:
            config[args.example_name]['masks_file'] = args.masks_file
        if not os.path.exists(config[args.example_name]['masks_file']):
            msg = '{} does not exist! You should specify masks_file correctly, ' \
                  'or use default one which is generated by model_prune_torch.py'
            raise RuntimeError(msg.format(config[args.example_name]['masks_file']))
        model_inference(config[args.example_name])
QuanluZhang's avatar
QuanluZhang committed
94
    else:
95
96
97
98
        model_inference(config['fpgm'])
        model_inference(config['slim'])
        model_inference(config['l1filter'])
        model_inference(config['apoz'])