def _create_estimator(spec, config, model_dir,
num_train_images, num_sample_images=None):
"""Creates the TPUEstimator object."""
# Estimator will save a checkpoint at the end of every train() call. Disable
# automatic checkpoints by setting the time interval between checkpoints to
# a very large value.
run_config = tf.contrib.tpu.RunConfig(
model_dir=model_dir,
keep_checkpoint_max=3, # Keeps ckpt at start, halfway, and end
save_checkpoints_secs=2**30,
tpu_config=tf.contrib.tpu.TPUConfig(
iterations_per_loop=config['tpu_iterations_per_loop'],
num_shards=config['tpu_num_shards']))
# This is a hack to allow PREDICT on a fixed batch on TPU. By replicating the
# batch by the number of shards, this ensures each TPU core operates on the
# entire fixed batch.
if num_sample_images and config['use_tpu']:
num_sample_images *= config['tpu_num_shards']
estimator = tf.contrib.tpu.TPUEstimator(
use_tpu=config['use_tpu'],
model_fn=model_builder.build_model_fn(
spec, config, num_train_images),
config=run_config,
train_batch_size=config['batch_size'],
eval_batch_size=config['batch_size'],
predict_batch_size=num_sample_images)
return estimator
def keytotuple(key):
cur_network_data = nb.get_metrics_from_hash(key)
#print(cur_network_data[0])
#print(cur_network_data[0].keys())
model = model_spec.ModelSpec(cur_network_data[0]['module_adjacency'],
cur_network_data[0]['module_operations'])
model_fn = model_builder.build_model_fn(model, cfg, 60000)
if os.path.exists('empty'):
shutil.rmtree('empty')
run_cfg = tf.contrib.tpu.RunConfig(
model_dir='empty',
keep_checkpoint_max=3, # Keeps ckpt at start, halfway, and end
save_checkpoints_secs=2**30)
#tpu_config=tf.contrib.tpu.TPUConfig(
# iterations_per_loop=cfg['tpu_iterations_per_loop'],
# num_shards=cfg['tpu_num_shards']))
#estimator = tf.contrib.tpu.TPUEstimator(model_fn, config=run_cfg,
# train_batch_size=cfg['batch_size'],
# eval_batch_size=cfg['batch_size'],
# predict_batch_size=cfg['batch_size'],
# use_tpu=False)#, params=cfg)
estimator = tf.estimator.Estimator(model_fn, config=run_cfg, params=cfg)
print(estimator)
#dummy_input = np.zeros((1, 224, 224, 3))
#dummy_label = np.zeros((1, 100))
#dummy_label[0] = 1
input_train = cifar.CIFARInput('train', cfg)
print(cfg['batch_size'])
#input_fn = tf.estimator.inputs.numpy_input_fn(x={"x": dummy_input}, y=dummy_label, shuffle=True)
#estimator.train(input_fn)
#estimator.train(input_fn=input_train.input_fn,
# max_steps=1)
#print(tf.get_default_graph().as_graph_def())
with tf.Graph().as_default() as g:
features = tf.placeholder(tf.float32, [cfg['batch_size'], 32, 32, 3])
labels = tf.placeholder(tf.int32, [cfg['batch_size']])
_ = model_fn(features,
labels,
mode=tf.estimator.ModeKeys.TRAIN,
params=cfg)
with tf.Session() as sess:
run_meta = tf.RunMetadata()
opts = tf.profiler.ProfileOptionBuilder.float_operation()
flops = tf.profiler.profile(g,
run_meta=run_meta,
cmd='op',
options=opts)
n_flops = flops.total_float_ops
print(n_flops)
#print(sess.graph.as_graph_def())
training_time_sum = 0.0
acc_sum = 0.0
params = cur_network_data[0]['trainable_parameters']
count = 0
for item in cur_network_data[1][108]:
count += 1
training_time_sum += item['final_training_time']
acc_sum += item['final_test_accuracy']
training_time = training_time_sum / count
acc = acc_sum / count
return (params, training_time, acc, n_flops)
def prepare_kd_dataset(spec, config, model_path, dataset_files,
new_dataset_path, trainset_part_percentage):
if config['dataset'] == 'mnist':
_dummy_imput_fn = _dummy_imput_fn_mnist
elif config['cifar']:
_dummy_imput_fn = _dummy_imput_fn_cifar
else:
raise Exception('Unsupported config[\'dataset\'] = {}. '
'Supported only cifar of mnist'.format(
config['dataset']))
Path(new_dataset_path).mkdir(parents=True, exist_ok=True)
for filename in dataset_files:
raw_dataset = tf.data.TFRecordDataset([filename])
params = {'file': filename, 'use_KD': False}
estimator = tf.contrib.tpu.TPUEstimator(
use_tpu=False,
model_fn=model_builder.build_model_fn(spec, config, None),
config=tf.contrib.tpu.RunConfig(model_dir=model_path),
params=params,
train_batch_size=config['batch_size'],
eval_batch_size=config['batch_size'],
predict_batch_size=100)
est_preds = estimator.predict(input_fn=_dummy_imput_fn,
yield_single_examples=False)
all_pred_logits_aug = []
for preds in est_preds:
all_pred_logits_aug.append(preds['logits'])
if len(all_pred_logits_aug) == 0:
logging.error(filename, ": all_pred_logits_aug is empty")
all_pred_logits_aug = np.vstack(all_pred_logits_aug)
filename = Path(filename)
name_postfix = '_KD'
if trainset_part_percentage != 100:
name_postfix += '_' + str(trainset_part_percentage)
out_file = filename.with_name(filename.stem + name_postfix)
out_file = out_file.with_suffix(".tfrecords")
out_file = Path(new_dataset_path, out_file.name)
filename = str(filename)
with tf.io.TFRecordWriter(str(out_file)) as record_writer:
for i, raw_record in enumerate(raw_dataset):
if i >= 10000 * (trainset_part_percentage / 100.0):
break
example = tf.train.Example()
example.ParseFromString(raw_record.numpy())
img = example.features.feature['image']
label = example.features.feature['label']
if len(all_pred_logits_aug) <= i:
logging.error(
"all_pred_logits_aug is not anought ({}); i={}".format(
len(all_pred_logits_aug), i))
continue
preds = all_pred_logits_aug[i]
new_label = np.hstack((label.int64_list.value[0], preds))
feat_preds = tf.train.Feature(float_list=tf.train.FloatList(
value=new_label))
example = tf.train.Example(features=tf.train.Features(
feature={
'image': img,
'label': feat_preds
}))
record_writer.write(example.SerializeToString())
logging.info("{} stored".format(out_file))