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