forked from eliberis/uNAS
-
Notifications
You must be signed in to change notification settings - Fork 0
/
metrics_file_ntk_rn.py
241 lines (199 loc) · 10.5 KB
/
metrics_file_ntk_rn.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
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
#################### uNAS model_trainer ####################
import logging
from typing import Optional
import tensorflow as tf
from config import TrainingConfig
from pruning import DPFPruning
from utils import debug_mode
#################### save_dataset save_model wait_metrics####################
import numpy as np
import tensorflow as tf
import tf2onnx
import onnx
import pickle
import time
import os
import pdb
import gc
#################### GPU ####################
# GPU mem issue
#config = tf.compat.v1.ConfigProto(gpu_options=tf.compat.v1.GPUOptions(allow_growth=True))
#sess = tf.compat.v1.Session(config=config)
class ModelMetricsFile:
"""Keras models according to the specified config."""
def __init__(self, trainer):
self.trainer = trainer
self.save_path = './tmp/metrics/ntk_rn'
def get_metrics(self, model, model_rn, num_batch, num_networks):
dataset = self.trainer.dataset
input_shape = self.trainer.dataset.input_shape
num_classes = self.trainer.dataset.num_classes
batch_size = self.trainer.config.batch_size
save_path = self.save_path
#save_dataset
save_path = self.save_path
dataset = self.trainer.dataset
batch_size = self.trainer.config.batch_size
num_batch = num_batch
#save_model
save_path = self.save_path
input_shape = self.trainer.dataset.input_shape
model = model
#save_model_rn
save_path = self.save_path
input_shape = self.trainer.dataset.input_shape
model_rn = model_rn
#wait_metrics
save_path = self.save_path
num_classes = self.trainer.dataset.num_classes
num_batch = num_batch
num_networks = num_networks
if not os.path.exists(save_path):
os.makedirs(save_path)
print (f"save_path dir:{save_path}")
else:
print (f"save_path is already exist:{save_path}")
def save_dataset(save_path, dataset, batch_size, num_batch, timestamp):
#################### dataset ####################
#check loader
train_loader_save_path = f'{save_path}/train_loader.pickle'
val_loader_save_path = f'{save_path}/val_loader.pickle'
if os.path.isfile(train_loader_save_path) and os.path.isfile(val_loader_save_path):
print (f"train_loader_save_path is already exist:{train_loader_save_path}")
print (f"val_loader_save_path is already exist:{val_loader_save_path}")
else:
print (f"generate train_loader_save_path :{train_loader_save_path}")
print (f"generate val_loader_save_path :{val_loader_save_path}")
# from uNAS dataset by tf
train = dataset.train_dataset() \
.shuffle(100000) \
.batch(batch_size) \
.prefetch(tf.data.experimental.AUTOTUNE)
# <PrefetchDataset shapes: ((None, 32, 32, 3), (None, 1)), types: (tf.float64, tf.uint8)>
val = dataset.validation_dataset() \
.batch(batch_size) \
.prefetch(tf.data.experimental.AUTOTUNE)
# <PrefetchDataset shapes: ((None, 32, 32, 3), (None, 1)), types: (tf.float64, tf.uint8)>
# for get ntk loader input
train_loader = []
val_loader = []
for i in range(num_batch):
# list(train.as_numpy_iterator())[0][0].shape = (128, 32, 32, 3)
train_input = list(train.as_numpy_iterator())[0][0]
# list(train.as_numpy_iterator())[0][1].shape = (128, 1)?
train_target = list(train.as_numpy_iterator())[0][1]
# targets.shape = torch.Size([128])
train_target = train_target.reshape((-1,))
# one_hot is only applicable to index tensor
train_target = train_target.astype(np.int64)
train_loader.append((train_input,train_target))
# for get ntk val_loader input
# list(val.as_numpy_iterator())[0][0].shape = (128, 32, 32, 3)
val_input = list(val.as_numpy_iterator())[0][0]
# list(val.as_numpy_iterator())[0][1].shape = (128, 1)?
val_target = list(val.as_numpy_iterator())[0][1]
# targets.shape = torch.Size([128])
val_target = val_target.reshape((-1,))
# one_hot is only applicable to index tensor
val_target = val_target.astype(np.int64)
val_loader.append((val_input,val_target))
# save loader as train_loader.pickle/val_loader.pickle
train_loader_save_path = f'{save_path}/train_loader.pickle'
with open(train_loader_save_path, 'wb') as f:
pickle.dump(train_loader, f)
val_loader_save_path = f'{save_path}/val_loader.pickle'
with open(val_loader_save_path, 'wb') as f:
pickle.dump(val_loader, f)
#clear the parameter
del train, val
del train_input, train_target
del val_input, val_target
gc.collect()
def save_model(save_path, input_shape, model: tf.keras.Model, timestamp):
#################### model ####################
# (load) model
keras_model = model
# input_shape like (None, 32, 32, 3)
input_shape = (None,) + input_shape
print (f'input_shape={input_shape}')
# tensorflow-onnx(維度可從dataset獲取)
keras_model_spec = (tf.TensorSpec(input_shape, tf.float32, name="input"),)
model_proto, external_tensor_storage = tf2onnx.convert.from_keras(keras_model,
input_signature=keras_model_spec, opset=None, custom_ops=None,
custom_op_handlers=None, custom_rewriter=None,
inputs_as_nchw=None, extra_opset=None, shape_override=None,
target=None, large_model=False, output_path=None)
onnx_model = model_proto
# Save the ONNX model
onnx_model_path = f'{save_path}/model_{timestamp}.onnx'
onnx.save(onnx_model, onnx_model_path)
#clear the parameter
del onnx_model, model_proto, external_tensor_storage, keras_model_spec, keras_model, model
gc.collect()
def save_model_rn(save_path, input_shape, model_rn, timestamp):
#################### model ####################
# (load) model
keras_model = model_rn
# input_shape like (1000, 3, 3, 1)
input_shape = (None, 2, 2, 1)
print (f'input_shape={input_shape}')
# tensorflow-onnx(維度可從dataset獲取)
keras_model_spec = (tf.TensorSpec(input_shape, tf.float32, name="input"),)
model_proto, external_tensor_storage = tf2onnx.convert.from_keras(keras_model,
input_signature=keras_model_spec, opset=None, custom_ops=None,
custom_op_handlers=None, custom_rewriter=None,
inputs_as_nchw=None, extra_opset=None, shape_override=None,
target=None, large_model=False, output_path=None)
onnx_model = model_proto
# Save the ONNX model
onnx_model_path = f'{save_path}/model_rn_{timestamp}.onnx'
onnx.save(onnx_model, onnx_model_path)
#clear the parameter
del onnx_model, model_proto, external_tensor_storage, keras_model_spec, keras_model, model_rn
gc.collect()
def wait_metrics(save_path, num_classes, num_batch, num_networks, timestamp):
# input_finish_info
print (f'num_batch={num_batch}')
print (f'num_classes={num_classes}')
print (f'num_networks={num_networks}')
print (f'timestamp={timestamp}')
input_finish_info = {"num_batch":num_batch, "num_classes":num_classes, "num_networks":num_networks,"timestamp":timestamp}
# save input_finish_info as input_finish_info.pickle
input_finish_info_save_path = f'{save_path}/input_finish_info_{timestamp}.pickle'
with open(input_finish_info_save_path, 'wb') as f:
pickle.dump(input_finish_info, f)
#check metrics_finish_info.pickle exsit
metrics_finish_info_save_path = f'{save_path}/metrics_finish_info_{timestamp}.pickle'
while not os.path.isfile(metrics_finish_info_save_path):
print (f'wait for metrics_finish_info_{timestamp}')
time.sleep(5)
# check the timestamp
while True:
time.sleep(5)
# load metrics_finish_info.pickle
with open(metrics_finish_info_save_path, 'rb') as f:
metrics_finish_info = pickle.load(f)
if timestamp != metrics_finish_info["timestamp"]:
print (f'the timestamp of metrics_finish_info is not right, wait for metrics_finish_info_{timestamp}.')
time.sleep(5)
elif timestamp == metrics_finish_info["timestamp"]:
print (f'find metrics_finish_info_{timestamp}')
break
time.sleep(5)
ntks = metrics_finish_info['ntks']
rns = metrics_finish_info['rns']
# del
os.remove(metrics_finish_info_save_path)
return ntks,rns
# timestamp
timestamp = "{:}".format(time.strftime('%h-%d-%C_%H-%M-%s', time.localtime(time.time())))
pid = os.getpid()
timestamp = f"{timestamp}_{pid}"
# save dataset
save_dataset(save_path=save_path , dataset=dataset, batch_size=batch_size, num_batch=num_batch, timestamp=timestamp)
# save model
save_model(save_path=save_path, input_shape=input_shape, model=model, timestamp=timestamp)
save_model_rn(save_path=save_path, input_shape=input_shape, model_rn=model_rn, timestamp=timestamp)
# wait ntk
ntks, rns = wait_metrics(save_path, num_classes=num_classes, num_batch=num_batch, num_networks=num_networks, timestamp=timestamp)
return ntks, rns