-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathgate_visualization.py
114 lines (90 loc) · 3.22 KB
/
gate_visualization.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
import torch
from dfme.train import train, test
from dfme.utils import get_dynamic_connections, get_dataloader
from models.resnet import *
from models.gated_resnet import DynamicResNet18_8x
from models.gan import GeneratorA
class ARGS:
epoch_iter = 50
g_iter = 1
s_iter = 5
num_classes = 10
dataset = 'cifar10'
data_root = './data/'
image_size = 32
image_channels = 3
batch_size = 256
student_lr = 0.1
generator_lr = 1e-4
weight_decay = 5e-4
steps = [0.1, 0.3, 0.5]
momentum = 0.9
random_noise_size = 256
grad_epsilon = 1e-3
grad_m = 1
current_budget = 0
cost_per_iteration = 0 # will be calculated within the program
query_budget = 10e6
seed = 4
generator_activation = torch.tanh
teacher = resnet34()
teacher_name = 'resnet34'
student = DynamicResNet18_8x(num_classes)
student_name = 'dynamic_resnet_18'
generator = GeneratorA(nz=random_noise_size, nc=image_channels, img_size=image_size, activation=generator_activation)
teacher_load_path = './checkpoints/teacher/resnet34.pt'
student_load_path = './checkpoints/student/dynamic_resnet_18_resnet34.pt' # fill after saving atleast once
generator_load_path = './checkpoints/generator/dynamic_resnet_18_resnet34.pt' # fill after saving atleast once
student_save_folder = './checkpoints/student/'
generator_save_folder = './checkpoints/generator'
args = ARGS()
acc = {}
f = open('results.txt', "a")
args.student_load_path = ''
args.generator_load_path = ''
print('\n\n\n')
# VISUALIZING DYNAMIC GATES
torch.cuda.empty_cache()
device=torch.device('cuda')
train_loader, test_loader = get_dataloader(args)
student = args.student.to(device)
student.load_state_dict(torch.load('./checkpoints/student/dynamic_resnet_18_resnet18.pt'))
test(student, device, test_loader=test_loader)
for i, (X, y) in enumerate(train_loader):
print(get_dynamic_connections(student, X[0].to(device)))
if (i == 5):
break
print('\n\n\n')
torch.cuda.empty_cache()
device=torch.device('cuda')
train_loader, test_loader = get_dataloader(args)
student = args.student.to(device)
student.load_state_dict(torch.load('./checkpoints/student/dynamic_resnet_18_resnet34.pt'))
test(student, device, test_loader=test_loader)
for i, (X, y) in enumerate(train_loader):
print(get_dynamic_connections(student, X[0].to(device)))
if (i == 5):
break
print('\n\n\n')
torch.cuda.empty_cache()
device=torch.device('cuda')
train_loader, test_loader = get_dataloader(args)
student = args.student.to(device)
student.load_state_dict(torch.load('./checkpoints/student/dynamic_resnet_18_resnet50.pt'))
test(student, device, test_loader=test_loader)
for i, (X, y) in enumerate(train_loader):
print(get_dynamic_connections(student, X[0].to(device)))
if (i == 5):
break
print('\n\n\n')
torch.cuda.empty_cache()
device=torch.device('cuda')
train_loader, test_loader = get_dataloader(args)
student = args.student.to(device)
student.load_state_dict(torch.load('./checkpoints/student/dynamic_resnet_18_googlenet.pt'))
test(student, device, test_loader=test_loader)
for i, (X, y) in enumerate(train_loader):
print(get_dynamic_connections(student, X[0].to(device)))
if (i == 5):
break
print('\n\n\n')