def __calc_rlout_reward(self, prune_ratios):
"""Calculate the reward of the current roll-out.
Args:
* prune_ratios: list of pruning ratios
Returns:
* reward: reward of the current roll-out
"""
# initialize the weight sparsified network with given pruning ratios
self.sess_train.run(self.pr_assign_op,
feed_dict={self.pr_all: prune_ratios})
self.sess_train.run([self.init_op, self.rg_init_op, self.ft_init_op])
if FLAGS.enbl_multi_gpu:
self.sess_train.run(self.bcast_op)
# evaluate the network before re-training
if is_primary_worker('global'):
loss_pre, metrics_pre = self.__calc_loss_n_metrics()
assert 'accuracy' in metrics_pre or 'acc_top5' in metrics_pre, \
'either <accuracy> or <acc_top5> must be evaluated and returned'
# re-train the network with layerwise regression & network fine-tuning
self.__retrain_network()
# evaluate the network after re-training
if is_primary_worker('global'):
loss_post, metrics_post = self.__calc_loss_n_metrics()
assert 'accuracy' in metrics_post or 'acc_top5' in metrics_post, \
'either <accuracy> or <acc_top5> must be evaluated and returned'
# evaluate the weight sparsified network
reward = None
if is_primary_worker('global'):
if 'accuracy' in metrics_post:
reward_pre = self.rl_helper.calc_reward(
metrics_pre['accuracy'])
reward = self.rl_helper.calc_reward(metrics_post['accuracy'])
elif 'acc_top5' in metrics_post:
reward_pre = self.rl_helper.calc_reward(
metrics_pre['acc_top5'])
reward = self.rl_helper.calc_reward(metrics_post['acc_top5'])
prune_ratio = self.rl_helper.calc_overall_prune_ratio()
metrics_diff = ' | '.join([
'%s: %.4f -> %.4f' % (key, metrics_pre[key], metrics_post[key])
for key in metrics_post
])
tf.logging.info(
'loss: %.4e -> %.4e | %s | reward: %.4f -> %.4f | prune_ratio = %.4f'
% (loss_pre, loss_post, metrics_diff, reward_pre, reward,
prune_ratio))
return reward
def dump_n_eval(self, outputs, action):
"""Dump the model's outputs to files and evaluate."""
if not is_primary_worker('global'):
return
if action == 'init':
if os.path.exists(FLAGS.outputs_dump_dir):
shutil.rmtree(FLAGS.outputs_dump_dir)
os.mkdir(FLAGS.outputs_dump_dir)
elif action == 'dump':
filename = outputs['predictions']['filename'][0].decode('utf8')[:-4]
shape = outputs['predictions']['shape'][0]
for idx_cls in range(1, FLAGS.nb_classes):
with open(os.path.join(FLAGS.outputs_dump_dir, 'results_%d.txt' % idx_cls), 'a') as o_file:
scores = outputs['predictions']['scores_%d' % idx_cls][0]
bboxes = outputs['predictions']['bboxes_%d' % idx_cls][0]
bboxes[:, 0] = (bboxes[:, 0] * shape[0]).astype(np.int32, copy=False) + 1
bboxes[:, 1] = (bboxes[:, 1] * shape[1]).astype(np.int32, copy=False) + 1
bboxes[:, 2] = (bboxes[:, 2] * shape[0]).astype(np.int32, copy=False) + 1
bboxes[:, 3] = (bboxes[:, 3] * shape[1]).astype(np.int32, copy=False) + 1
for idx_bbox in range(bboxes.shape[0]):
bbox = bboxes[idx_bbox][:]
if bbox[2] > bbox[0] and bbox[3] > bbox[1]:
o_file.write('%s %.3f %.1f %.1f %.1f %.1f\n'
% (filename, scores[idx_bbox], bbox[1], bbox[0], bbox[3], bbox[2]))
elif action == 'eval':
do_python_eval(os.path.join(self.dataset_eval.data_dir, 'test'), FLAGS.outputs_dump_dir)
else:
raise ValueError('unrecognized action in dump_n_eval(): ' + action)
def dump_n_eval(self, outputs, action):
"""Dump the model's outputs to files and evaluate."""
if not is_primary_worker('global'):
return
if action == 'init':
if os.path.exists(FLAGS.outputs_dump_dir):
shutil.rmtree(FLAGS.outputs_dump_dir)
os.mkdir(FLAGS.outputs_dump_dir)
elif action == 'dump':
image_id = outputs['predictions']['image_id'][0]
shape = outputs['predictions']['shape'][0]
for idx_cls in range(1, FLAGS.nb_classes):
with open(
os.path.join(FLAGS.outputs_dump_dir,
'results_%d.txt' % idx_cls),
'a') as o_file:
scores = outputs['predictions']['scores_%d' % idx_cls][0]
bboxes = outputs['predictions']['bboxes_%d' % idx_cls][0]
bboxes[:, 0] = (bboxes[:, 0] * shape[0])
bboxes[:, 1] = (bboxes[:, 1] * shape[1])
bboxes[:, 2] = (bboxes[:, 2] * shape[0])
bboxes[:, 3] = (bboxes[:, 3] * shape[1])
for idx_bbox in range(bboxes.shape[0]):
bbox = bboxes[idx_bbox][:]
if bbox[2] > bbox[0] and bbox[3] > bbox[1]:
# data format convert
x = float(format(bbox[1], '.2f'))
y = float(format(bbox[0], '.2f'))
delta_x = float(format(bbox[3] - bbox[1], '.2f'))
delta_y = float(format(bbox[2] - bbox[0], '.2f'))
score = float(format(scores[idx_bbox], '.3f'))
# image_info = dict()
# image_info['image_id'] = int(image_id)
# image_info['category_id'] = idx_cls
# image_info['bbox'] = [x, y, delta_x, delta_y]
# image_info['score'] = score
#
# o_file.write(str(image_info)+',')
# [{"image_id":42,"category_id":18,"bbox":[258.15,41.29,348.26,243.78],"score":0.236},
o_file.write('%d %d %.2f %.2f %.2f %.2f %.3f\n' %
(int(image_id), idx_cls, x, y,
delta_x, delta_y, score))
elif action == 'eval':
do_python_eval(FLAGS.data_dir_local, FLAGS.outputs_dump_dir)
else:
raise ValueError('unrecognized action in dump_n_eval(): ' + action)
def dump_n_eval(self, outputs, action):
"""Dump the model's outputs to files and evaluate."""
if not is_primary_worker('global'):
return
if action == 'init':
if os.path.exists(FLAGS.outputs_dump_dir):
shutil.rmtree(FLAGS.outputs_dump_dir)
os.mkdir(FLAGS.outputs_dump_dir)
elif action == 'dump':
filename = outputs['predictions']['filename'][0].decode('utf8')[:-4]
raw_shape = outputs['predictions']['shape'][0]
resized_shape= outputs['predictions']['resized_shape']
detected_boxes = outputs['predictions']['detected_boxes']
detected_scores = outputs['predictions']['detected_scores']
detected_categories = outputs['predictions']['detected_categories']
raw_h, raw_w = raw_shape[0], raw_shape[1]
resized_h, resized_w = resized_shape[1], resized_shape[2]
xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
detected_boxes[:, 2], detected_boxes[:, 3]
xmin = xmin * raw_w / resized_w
xmax = xmax * raw_w / resized_w
ymin = ymin * raw_h / resized_h
ymax = ymax * raw_h / resized_h
boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
dets = np.hstack((detected_categories.reshape(-1, 1),
detected_scores.reshape(-1, 1),
boxes))
for cls_id in range(1, FLAGS.nb_classes):
with open(os.path.join(FLAGS.outputs_dump_dir, 'results_%d.txt' % cls_id), 'a') as o_file:
this_cls_detections = dets[dets[:, 0] == cls_id]
if this_cls_detections.shape[0] == 0:
continue # this cls has none detections in this img
for a_det in this_cls_detections:
o_file.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.
format(filename, a_det[1],
a_det[2], a_det[3],
a_det[4], a_det[5])) # that is [img_name, score, xmin, ymin, xmax, ymax]
elif action == 'eval':
do_python_eval(os.path.join(self.dataset_eval.data_dir, 'test'), FLAGS.outputs_dump_dir)
else:
raise ValueError('unrecognized action in dump_n_eval(): ' + action)
def __init__(self, sm_writer, model_helper, mpi_comm):
"""Constructor function.
Args:
* sm_writer: TensorFlow's summary writer
* model_helper: model helper with definitions of model & dataset
* mpi_comm: MPI communication object
"""
# initialize a full-precision model
self.model_scope = 'distilled_model' # to distinguish from models created by other learners
enbl_dst = False # disable the distillation loss for teacher model
from learners.full_precision.learner import FullPrecLearner
self.learner = FullPrecLearner(sm_writer, model_helper,
self.model_scope, enbl_dst)
# initialize a model for training with the distillation loss
if is_primary_worker('local'):
self.__initialize()
if FLAGS.enbl_multi_gpu:
mpi_comm.Barrier()
def run(self):
"""Run the optimizer to obtain pruning ratios for all maskable variables.
Returns:
* var_names_n_prune_ratios: list of variable name & pruning ratio pairs
"""
# obtain a list of (variable name, pruning ratio) tuples
if FLAGS.ws_prune_ratio_prtl == 'uniform':
var_names_n_prune_ratios = self.__calc_uniform_prune_ratios()
elif FLAGS.ws_prune_ratio_prtl == 'heurist':
var_names_n_prune_ratios = self.__calc_heurist_prune_ratios()
elif FLAGS.ws_prune_ratio_prtl == 'optimal':
var_names_n_prune_ratios = self.__calc_optimal_prune_ratios()
# display the pruning ratio for each maskable variable
if is_primary_worker('global'):
for var_name, prune_ratio in var_names_n_prune_ratios:
tf.logging.info('%s: %f' % (var_name, prune_ratio))
return var_names_n_prune_ratios
def __calc_optimal_prune_ratios(self):
"""Calculate pruning ratios using the 'optimal' protocol.
Returns:
* var_names_n_prune_ratios: list of variable name & pruning ratio pairs
"""
# restore the full-precision model from checkpoint files
save_path = tf.train.latest_checkpoint(
os.path.dirname(self.save_path_full))
self.saver_full.restore(self.sess_train, save_path)
# train an RL agent through multiple roll-outs
if is_primary_worker('global'):
self.agent.init()
reward_best = np.NINF # pylint: disable=no-member
prune_ratios_best = None
file_path_prune_ratios = './ws.prune.ratios'
file_path_reward = './ws.reward'
for idx_rlout in range(FLAGS.ws_nb_rlouts):
# compute actions & pruning ratios
if is_primary_worker('global'):
tf.logging.info('starting %d-th roll-out' % idx_rlout)
prune_ratios, states_n_actions = self.__calc_rlout_actions()
save_vals_to_file(prune_ratios, file_path_prune_ratios)
if FLAGS.enbl_multi_gpu:
self.mpi_comm.Barrier()
prune_ratios = restore_vals_from_file(file_path_prune_ratios)
# fine-tune the weight sparsified network to compute the reward
reward = self.__calc_rlout_reward(prune_ratios)
if is_primary_worker('global'):
save_vals_to_file(np.array([reward]), file_path_reward)
if FLAGS.enbl_multi_gpu:
self.mpi_comm.Barrier()
reward = restore_vals_from_file(file_path_reward)[0]
# update the baseline function in DDPG
if is_primary_worker('global'):
rewards = reward * np.ones(len(self.vars_full['maskable']))
self.agent.finalize_rlout(rewards)
# record transitions to train the RL agent
if is_primary_worker('global'):
self.__record_rlout_transitions(states_n_actions, reward)
# record the best combination of pruning ratios
if reward_best < reward:
if is_primary_worker('global'):
tf.logging.info('best reward updated: %.4f -> %.4f' %
(reward_best, reward))
tf.logging.info('optimal pruning ratios: ' + ' '.join([
'%.2f' % prune_ratio for prune_ratio in prune_ratios[:]
]))
reward_best = reward
prune_ratios_best = np.copy(prune_ratios)
# setup the optimal pruning ratios
var_names_n_prune_ratios = []
for idx, var_full in enumerate(self.vars_full['maskable']):
var_names_n_prune_ratios += [(var_full.name,
prune_ratios_best[idx])]
return var_names_n_prune_ratios
def __build_train(self, model_helper): # pylint: disable=too-many-locals
"""Build the training graph for the 'optimal' protocol.
Args:
* model_helper: model helper with definitions of model & dataset
"""
with tf.Graph().as_default():
# create a TF session for the current graph
config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(
mgw.local_rank() if FLAGS.enbl_multi_gpu else 0) # pylint: disable=no-member
sess = tf.Session(config=config)
# data input pipeline
with tf.variable_scope(self.data_scope):
iterator, __ = model_helper.build_dataset_train(
enbl_trn_val_split=True)
images, labels = iterator.get_next()
# model definition - full-precision network
with tf.variable_scope(self.model_scope_full):
logits = model_helper.forward_eval(
images) # DO NOT USE forward_train() HERE!!!
self.vars_full = get_vars_by_scope(self.model_scope_full)
self.saver_full = tf.train.Saver(self.vars_full['all'])
self.save_path_full = FLAGS.save_path
# model definition - weight sparsified network
with tf.variable_scope(self.model_scope_prnd):
# forward pass & variables' saver
logits = model_helper.forward_eval(
images) # DO NOT USE forward_train() HERE!!!
self.vars_prnd = get_vars_by_scope(self.model_scope_prnd)
self.maskable_var_names = [
var.name for var in self.vars_prnd['maskable']
]
loss, __ = model_helper.calc_loss(labels, logits,
self.vars_prnd['trainable'])
self.saver_prnd_train = tf.train.Saver(self.vars_prnd['all'])
self.save_path_prnd = FLAGS.save_path.replace(
'models', 'models_pruned')
# build masks for variable pruning
self.masks, self.pr_all, self.pr_assign_op = self.__build_masks(
)
# create operations for initializing the weight sparsified network
init_ops = []
for var_full, var_prnd in zip(self.vars_full['all'],
self.vars_prnd['all']):
if var_full not in self.vars_full['maskable']:
init_ops += [var_prnd.assign(var_full)]
else:
idx = self.vars_full['maskable'].index(var_full)
init_ops += [var_prnd.assign(var_full * self.masks[idx])]
self.init_op = tf.group(init_ops)
# build operations for layerwise regression & network fine-tuning
self.rg_init_op, self.rg_train_ops = self.__build_layer_rg_ops()
self.ft_init_op, self.ft_train_op = self.__build_network_ft_ops(
loss)
if FLAGS.enbl_multi_gpu:
self.bcast_op = mgw.broadcast_global_variables(0)
# create RL helper & agent on the primary worker
if is_primary_worker('global'):
self.rl_helper, self.agent = self.__build_rl_helper_n_agent(
sess)
# TF operations
self.sess_train = sess
