-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmain_query.py
159 lines (134 loc) · 6.25 KB
/
main_query.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
import os
import numpy as np
import argparse
import random
import tensorflow as tf
from data_generators.data_generator import DataGenerator
from models.maml import MAML
parser = argparse.ArgumentParser()
parser.add_argument('data_path', metavar='DATA',
help='path to data')
parser.add_argument('ckpt_name', metavar='CKPT',
help='path to checkpoint')
parser.add_argument('-ds', '--data_source', default='imagenet', type=str,
help='data_source (imagenet or omniglot)')
parser.add_argument('-t', '--test', action='store_true', default=False,
help='set for test data, otherwise training data')
parser.add_argument('--multi', action='store_true', default=False,
help='set for multi-class problems, otherwise binary classification')
parser.add_argument('-l', '--train_lr', default=1e-4, type=float,
help='train_lr (default=1e-4)')
parser.add_argument('-p', '--pkl_file', default='filelist', type=str,
help='path to pickle file')
parser.add_argument('-cf', '--cluster_folder', default=None, type=str,
help='cluster folder w/o root (default=None)')
parser.add_argument('-cl', '--num_clusters', default=16, type=int,
help='# of clusters (default=16)')
parser.add_argument('-m', '--model_id', default="0", type=str,
help='model ID (default="0")')
# use kshot = 1, kquery = 15, nway = 5 for 5-way one-shot (multi-class)
parser.add_argument('--kshot', default=1, type=int,
help='# of shots per class (default=1)')
parser.add_argument('--kquery', default=1, type=int,
help='# of queries per class (default=1)')
parser.add_argument('--nway', default=51, type=int,
help='# of classes per problem (default=51)')
parser.add_argument('--metabatch', default=4, type=int,
help='meta batch-size for training (default=4)')
parser.add_argument('--steps', default=5, type=int,
help='# of gradient steps (default=5)')
parser.add_argument('--train_problems', default=40000, type=int,
help='# of training problems (default=40,000)')
parser.add_argument('--test_problems', default=10000, type=int,
help='# of test problems (default=10,000)')
parser.add_argument('-c', '--cuda_id', default="0", type=str,
help='cuda ID (default="0")')
args = parser.parse_args()
data_path = args.data_path
data_source = args.data_source
ckpt_name = args.ckpt_name
train_lr = args.train_lr
pkl_file = args.pkl_file
cluster_folder = args.cluster_folder
num_clusters = args.num_clusters
model_id = args.model_id
kshot = args.kshot
kquery = args.kquery
nway = args.nway
meta_batchsz = args.metabatch
steps = args.steps
train_problems = args.train_problems
test_problems = args.test_problems
cuda_id = args.cuda_id
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = cuda_id
def get_preds(model, multiclass, sess):
np.random.seed(1)
random.seed(1)
if args.test:
n_iterations = test_problems // meta_batchsz
else:
n_iterations = train_problems // meta_batchsz
query_acc = []
query_preds = []
for i in range(n_iterations):
ops = [model.test_query_accs, model.query_preds]
result = sess.run(ops)
print('Accuracy:', result[0])
query_acc.extend(list(result[0]))
query_preds.extend(result[1])
if multiclass:
# always do on clusters
np.save(os.path.join(data_path, 'CLUSTER_' + str(num_clusters),
'queryPredsMulticlass_' + pkl_file + '_cluster'
+ str(num_clusters) + '_' + model_id), query_preds)
else:
if cluster_folder is None:
np.save(os.path.join(data_path, 'WHOLE',
'queryPreds_' + pkl_file
+ '_model' + model_id), query_preds)
np.save(os.path.join(data_path, 'WHOLE',
'queryAcc_' + pkl_file
+ '_model' + model_id), query_acc)
else:
np.save(os.path.join(data_path, 'CLUSTER_' + str(num_clusters),
'queryPreds_' + pkl_file + '_cluster'
+ str(num_clusters) + '_' + model_id), query_preds)
np.save(os.path.join(data_path, 'CLUSTER_' + str(num_clusters),
'queryAcc_' + pkl_file + '_cluster'
+ str(num_clusters) + '_' + model_id), query_acc)
print('****DONE*****')
def main():
multiclass = args.multi
# kshot + kquery images per category, nway categories, meta_batchsz tasks.
db = DataGenerator(data_source, nway, kshot, kquery, meta_batchsz,
pkl_file, data_path, cluster_folder, multiclass, train_problems, test_problems)
image_tensor, label_tensor = db.make_data_tensor(training=True)
# get the tensors
support_x = tf.slice(image_tensor, [0, 0, 0], [-1, nway, -1], name='support_x')
query_x = tf.slice(image_tensor, [0, nway, 0], [-1, -1, -1], name='query_x')
support_y = tf.slice(label_tensor, [0, 0, 0], [-1, nway, -1], name='support_y')
query_y = tf.slice(label_tensor, [0, nway, 0], [-1, -1, -1], name='query_y')
model = MAML(data_source, 2, kshot, kquery, train_lr=train_lr)
model.build(support_x, support_y, query_x, query_y, steps, meta_batchsz, mode='testEach')
all_vars = filter(lambda x: 'meta_optim' not in x.name, tf.trainable_variables())
for p in all_vars:
print(p)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.InteractiveSession(config=config)
# tf.global_variables() to save moving_mean and moving variance of batch norm
# tf.trainable_variables() NOT include moving_mean and moving_variance.
saver = tf.train.Saver(tf.global_variables(), max_to_keep=3)
# initialize, under interative session
tf.global_variables_initializer().run()
tf.train.start_queue_runners()
if os.path.exists(os.path.join(ckpt_name, 'checkpoint')):
# alway load ckpt both train and test.
model_file = tf.train.latest_checkpoint(ckpt_name)
print("Restoring model weights from ", model_file)
saver.restore(sess, model_file)
get_preds(model, multiclass, sess)
if __name__=="__main__":
main()