-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathextract_qrcode_res18_cifar100.py
181 lines (144 loc) · 6.27 KB
/
extract_qrcode_res18_cifar100.py
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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
'''
load lottery tickets and evaluation
support datasets: cifar10, Fashionmnist, cifar100
'''
import os
import time
import random
import shutil
import argparse
import numpy as np
from copy import deepcopy
import matplotlib.pyplot as plt
import torch
import torch.optim
import torch.nn as nn
import torch.utils.data
import torch.nn.functional as F
import torchvision.models as models
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data.sampler import SubsetRandomSampler
from advertorch.utils import NormalizeByChannelMeanStd
from utils import *
from pruning_utils_2 import *
from pruning_utils_unprune import *
parser = argparse.ArgumentParser(description='PyTorch Evaluation Tickets')
##################################### general setting #################################################
parser.add_argument('--data', type=str, default='../../data', help='location of the data corpus')
parser.add_argument('--dataset', type=str, default='cifar10', help='dataset')
parser.add_argument('--arch', type=str, default='res18', help='model architecture')
parser.add_argument('--seed', default=None, type=int, help='random seed')
parser.add_argument('--save_dir', help='The directory used to save the trained models', default=None, type=str)
parser.add_argument('--gpu', type=int, default=0, help='gpu device id')
parser.add_argument('--save_model', action="store_true", help="whether saving model")
##################################### training setting #################################################
parser.add_argument('--optim', type=str, default='sgd', help='optimizer')
parser.add_argument('--batch_size', type=int, default=128, help='batch size')
parser.add_argument('--lr', default=0.1, type=float, help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
parser.add_argument('--weight_decay', default=1e-4, type=float, help='weight decay')
parser.add_argument('--epochs', default=182, type=int, help='number of total epochs to run')
parser.add_argument('--warmup', default=0, type=int, help='warm up epochs')
parser.add_argument('--print_freq', default=50, type=int, help='print frequency')
parser.add_argument('--decreasing_lr', default='91,136', help='decreasing strategy')
##################################### Pruning setting #################################################
parser.add_argument('--pretrained', default=None, type=str, help='pretrained weight for pt')
parser.add_argument('--mask_dir', default=None, type=str, help='mask direction for ticket')
parser.add_argument('--conv1', action="store_true", help="whether pruning&rewind conv1")
parser.add_argument('--fc', action="store_true", help="whether rewind fc")
parser.add_argument('--type', type=str, default=None, choices=['ewp', 'random_path', 'betweenness', 'hessian_abs', 'taylor1_abs','intgrads','identity'])
parser.add_argument('--add-back', action="store_true", help="add back weights")
parser.add_argument('--prune-type', type=str, choices=["lt", 'pt', 'st', 'mt', 'trained', 'transfer'])
parser.add_argument('--evaluate-p', type=float, default=0.00)
parser.add_argument('--evaluate-random', action='store_true')
parser.add_argument('--max-name', type=str)
parser.add_argument('--checkpoint', type=str)
best_sa = 0
def main():
global args, best_sa
args = parser.parse_args()
print(args)
print('*'*50)
print('conv1 included for prune and rewind: {}'.format(args.conv1))
print('fc included for rewind: {}'.format(args.fc))
print('*'*50)
os.makedirs(args.save_dir, exist_ok=True)
if args.seed:
setup_seed(args.seed)
# prepare dataset
model, train_loader, val_loader, test_loader = setup_model_dataset(args)
criterion = nn.CrossEntropyLoss()
state_dict = torch.load(args.checkpoint, map_location="cpu")['state_dict']
current_mask = extract_mask(state_dict)
print(current_mask.keys())
print(current_mask['layer3.1.conv1.weight_mask'].shape)
mask = current_mask['layer3.1.conv1.weight_mask'] > 0
prune_model_custom(model, current_mask, conv1=False)
check_sparsity(model, conv1=False)
try:
model.load_state_dict(state_dict)
except:
state_dict['normalize.mean'] = model.state_dict()['normalize.mean']
state_dict['normalize.std'] = model.state_dict()['normalize.std']
model.load_state_dict(state_dict)
validate(val_loader, model, criterion)
if args.evaluate_p > 0:
pruning_model(model, args.evaluate_p, random=args.evaluate_random)
check_sparsity(model, conv1=False)
state = model.state_dict()
mask = state[args.max_name + ".weight_mask"]
mask = mask.sum((2,3)) > 0
mask = mask.int().numpy()
h = 0
w = 33
mask = mask[h:h+29,w:w+29]
#plt.imshow(mask)
#plt.savefig(f'ownership/{args.arch}_{args.dataset}_qrcode_{args.evaluate_p}.png')
#assert False
#assert False
import qrcode
qr = qrcode.QRCode(
version=3,
error_correction=qrcode.constants.ERROR_CORRECT_H,
box_size=1,
border=0,
)
qr.add_data('signature')
qr.make(fit=True)
img = qr.make_image(fill_color="black", back_color="white")
code = np.array(img)
mask[0:9, 0:9] = code[0:9, 0:9]
mask[-9:,:9] = code[-9:,:9]
mask[:9,-9:] = code[:9,-9:]
mask[20:25, 20:25] = code[20:25, 20:25]
mask[-8, 4 * 3 + 9] = 1
mask[6] = code[6]
mask[:, 6] = code[:, 6]
print((mask != code).mean())
plt.imshow(mask)
plt.savefig(f'ownership/{args.arch}_{args.dataset}_qrcode_{args.evaluate_p}.png')
def save_checkpoint(state, is_SA_best, save_path, filename='checkpoint.pth.tar'):
filepath = os.path.join(save_path, filename)
torch.save(state, filepath)
if is_SA_best:
shutil.copyfile(filepath, os.path.join(save_path, 'model_SA_best.pth.tar'))
def validate(val_loader, model, criterion):
"""
Run evaluation
"""
# switch to evaluate mode
model.eval()
for i, (image, target) in enumerate(val_loader):
# compute output
with torch.no_grad():
output = model(image)
break
def setup_seed(seed):
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
if __name__ == '__main__':
main()