def test_prune_model_2_layers(self): """ Punning two layers with 0.5 comp-ratio in MNIST""" # create tf.compat.v1.Session and initialize the weights and biases with zeros config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True # create session with graph sess = tf.compat.v1.Session(graph=tf.Graph(), config=config) with sess.graph.as_default(): # by default, model will be constructed in default graph _ = mnist_tf_model.create_model(data_format='channels_last') sess.run(tf.compat.v1.global_variables_initializer()) # Create a layer database orig_layer_db = LayerDatabase(model=sess, input_shape=(1, 28, 28, 1), working_dir=None) conv1 = orig_layer_db.find_layer_by_name('conv2d/Conv2D') conv2 = orig_layer_db.find_layer_by_name('conv2d_1/Conv2D') layer_comp_ratio_list = [ LayerCompRatioPair(conv1, Decimal(0.5)), LayerCompRatioPair(conv2, Decimal(0.5)) ] spatial_svd_pruner = SpatialSvdPruner() comp_layer_db = spatial_svd_pruner.prune_model(orig_layer_db, layer_comp_ratio_list, CostMetric.mac, trainer=None) conv1_a = comp_layer_db.find_layer_by_name('conv2d_a/Conv2D') conv1_b = comp_layer_db.find_layer_by_name('conv2d_b/Conv2D') # Weights shape [kh, kw, Nic, Noc] self.assertEqual([5, 1, 1, 2], conv1_a.module.inputs[1].get_shape().as_list()) self.assertEqual([1, 5, 2, 32], conv1_b.module.inputs[1].get_shape().as_list()) conv2_a = comp_layer_db.find_layer_by_name('conv2d_1_a/Conv2D') conv2_b = comp_layer_db.find_layer_by_name('conv2d_1_b/Conv2D') self.assertEqual([5, 1, 32, 53], conv2_a.module.inputs[1].get_shape().as_list()) self.assertEqual([1, 5, 53, 64], conv2_b.module.inputs[1].get_shape().as_list()) for layer in comp_layer_db: print("Layer: " + layer.name) print(" Module: " + str(layer.module.name)) tf.compat.v1.reset_default_graph() sess.close() # delete temp directory shutil.rmtree(str('./temp_meta/'))
def test_layer_database_with_dynamic_shape(self): """ test layer database creation with different input shapes""" # create tf.compat.v1.Session and initialize the weights and biases with zeros config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True graph = tf.Graph() with graph.as_default(): # by default, model will be constructed in default graph input_placeholder = tf.compat.v1.placeholder(tf.float32, [None, None, None, 3], 'input') x = tf.keras.layers.Conv2D(8, (2, 2), padding='SAME')(input_placeholder) x = tf.keras.layers.BatchNormalization(momentum=.3, epsilon=.65)(x) x = tf.keras.layers.Conv2D(8, (1, 1), padding='SAME', activation=tf.nn.tanh)(x) x = tf.keras.layers.BatchNormalization(momentum=.4, epsilon=.25)(x) init = tf.compat.v1.global_variables_initializer() # create session with graph sess = tf.compat.v1.Session(graph=graph, config=config) sess.run(init) layer_db = LayerDatabase(model=sess, input_shape=(1, 224, 224, 3), working_dir=None, starting_ops=['input'], ending_ops=['batch_normalization_1/cond/Merge']) conv1_layer = layer_db.find_layer_by_name('conv2d/Conv2D') conv2_layer = layer_db.find_layer_by_name('conv2d_1/Conv2D') self.assertEqual(conv1_layer.output_shape, [1, 8, 224, 224]) self.assertEqual(conv2_layer.output_shape, [1, 8, 224, 224]) layer_db.destroy() # 2) try with different input shape # create another session with graph sess = tf.compat.v1.Session(graph=graph, config=config) sess.run(init) batch_size = 32 layer_db = LayerDatabase(model=sess, input_shape=(batch_size, 28, 28, 3), working_dir=None, starting_ops=['input'], ending_ops=['batch_normalization_1/cond/Merge']) conv1_layer = layer_db.find_layer_by_name('conv2d/Conv2D') conv2_layer = layer_db.find_layer_by_name('conv2d_1/Conv2D') self.assertEqual(conv1_layer.output_shape, [32, 8, 28, 28]) self.assertEqual(conv2_layer.output_shape, [32, 8, 28, 28]) layer_db.destroy()
def test_prune_model_tf_slim(self): """ Punning a model with tf slim api """ # create tf.compat.v1.Session and initialize the weights and biases with zeros config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True # create session with graph sess = tf.compat.v1.Session(graph=tf.Graph(), config=config) with sess.graph.as_default(): # by default, model will be constructed in default graph x = tf.compat.v1.placeholder(tf.float32, [1, 32, 32, 3]) _ = tf_slim_basic_model(x) sess.run(tf.compat.v1.global_variables_initializer()) conn_graph_orig = ConnectedGraph(sess.graph, ['Placeholder'], ['tf_slim_model/Softmax']) num_ops_orig = len(conn_graph_orig.get_all_ops()) # Create a layer database orig_layer_db = LayerDatabase(model=sess, input_shape=(1, 32, 32, 3), working_dir=None) conv1 = orig_layer_db.find_layer_by_name('Conv_1/Conv2D') conv1_bias = BiasUtils.get_bias_as_numpy_data(orig_layer_db.model, conv1.module) layer_comp_ratio_list = [LayerCompRatioPair(conv1, Decimal(0.5))] spatial_svd_pruner = SpatialSvdPruner() comp_layer_db = spatial_svd_pruner.prune_model(orig_layer_db, layer_comp_ratio_list, CostMetric.mac, trainer=None) # Check that svd added these ops _ = comp_layer_db.model.graph.get_operation_by_name('Conv_1_a/Conv2D') _ = comp_layer_db.model.graph.get_operation_by_name('Conv_1_b/Conv2D') conn_graph_new = ConnectedGraph(comp_layer_db.model.graph, ['Placeholder'], ['tf_slim_model/Softmax']) num_ops_new = len(conn_graph_new.get_all_ops()) self.assertEqual(num_ops_orig + 1, num_ops_new) bias_add_op = comp_layer_db.model.graph.get_operation_by_name( 'Conv_1_b/BiasAdd') conv_1_b_op = comp_layer_db.model.graph.get_operation_by_name( 'Conv_1_b/Conv2D') self.assertEqual( conn_graph_new._module_identifier.get_op_info(bias_add_op), conn_graph_new._module_identifier.get_op_info(conv_1_b_op)) self.assertTrue( np.array_equal( conv1_bias, BiasUtils.get_bias_as_numpy_data(comp_layer_db.model, conv_1_b_op)))
def test_prune_conv_no_bias(self): """ Test spatial svd on a conv layer with no bias """ # create tf.compat.v1.Session and initialize the weights and biases with zeros config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True # create session with graph sess = tf.compat.v1.Session(graph=tf.Graph(), config=config) with sess.graph.as_default(): # by default, model will be constructed in default graph inputs = tf.keras.Input(shape=( 32, 32, 3, )) x = tf.keras.layers.Conv2D(32, (3, 3), use_bias=False)(inputs) _ = tf.keras.layers.Flatten()(x) sess.run(tf.compat.v1.global_variables_initializer()) # Create a layer database orig_layer_db = LayerDatabase(model=sess, input_shape=(1, 32, 32, 3), working_dir=None) conv_op = orig_layer_db.find_layer_by_name('conv2d/Conv2D') layer_comp_ratio_list = [LayerCompRatioPair(conv_op, Decimal(0.5))] spatial_svd_pruner = SpatialSvdPruner() comp_layer_db = spatial_svd_pruner.prune_model(orig_layer_db, layer_comp_ratio_list, CostMetric.mac, trainer=None) # Check that svd added these ops _ = comp_layer_db.model.graph.get_operation_by_name('conv2d_a/Conv2D') conv2d_b_op = comp_layer_db.model.graph.get_operation_by_name( 'conv2d_b/Conv2D') reshape_op = comp_layer_db.model.graph.get_operation_by_name( 'flatten/Reshape') self.assertEqual(conv2d_b_op, reshape_op.inputs[0].op)
def _reconstruct_layers(self, layers_to_reconstruct: List[Layer], orig_layer_name_to_pruned_name_and_mask_dict: Dict[ str, Tuple[str, List[int]]], layer_db: LayerDatabase, comp_layer_db: LayerDatabase): """ Reconstruct weights and biases of layers in the layers_to_reconstruct list. :param layers_to_reconstruct: List of layers to reconstruct weights and biases of :param orig_layer_name_to_pruned_name_and_mask_dict: Dictionary mapping original layer names to most recent pruned op name and most recent output masks. :param layer_db: Original layer database :param comp_layer_db: Compressed layer database """ for layer in layers_to_reconstruct: # Get output mask of layer, that contains information about all channels winnowed since the start pruned_layer_name, output_mask = \ orig_layer_name_to_pruned_name_and_mask_dict.get(layer.name, (None, None)) assert pruned_layer_name is not None pruned_layer = comp_layer_db.find_layer_by_name(pruned_layer_name) self._data_subsample_and_reconstruction(layer, pruned_layer, output_mask, layer_db, comp_layer_db)
def test_per_layer_eval_scores(self): pruner = unittest.mock.MagicMock() eval_func = unittest.mock.MagicMock() # create tf.compat.v1.Session and initialize the weights and biases with zeros config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True # create session with graph sess = tf.compat.v1.Session(graph=tf.Graph(), config=config) with sess.graph.as_default(): # by default, model will be constructed in default graph _ = mnist_tf_model.create_model(data_format='channels_last') sess.run(tf.compat.v1.global_variables_initializer()) # Create a layer database layer_db = LayerDatabase(model=sess, input_shape=(1, 28, 28, 1), working_dir=None) layer1 = layer_db.find_layer_by_name('conv2d/Conv2D') layer_db.mark_picked_layers([layer1]) eval_func.side_effect = [90, 80, 70, 60, 50, 40, 30, 20, 10] url, process = start_bokeh_server_session(8006) bokeh_session = BokehServerSession(url=url, session_id="compression") # Instantiate child greedy_algo = comp_ratio_select.GreedyCompRatioSelectAlgo( layer_db=layer_db, pruner=pruner, cost_calculator=SpatialSvdCostCalculator(), eval_func=eval_func, eval_iterations=20, cost_metric=CostMetric.mac, target_comp_ratio=0.5, num_candidates=10, use_monotonic_fit=True, saved_eval_scores_dict=None, comp_ratio_rounding_algo=None, use_cuda=False, bokeh_session=bokeh_session) progress_bar = ProgressBar(1, "eval scores", "green", bokeh_session=bokeh_session) data_table = DataTable(num_columns=3, num_rows=1, column_names=[ '0.1', '0.2', '0.3', '0.4', '0.5', '0.6', '0.7', '0.8', '0.9' ], row_index_names=[layer1.name], bokeh_session=bokeh_session) pruner.prune_model.return_value = layer_db eval_dict = greedy_algo._compute_layerwise_eval_score_per_comp_ratio_candidate( data_table, progress_bar, layer1) self.assertEqual(90, eval_dict[Decimal('0.1')]) tf.compat.v1.reset_default_graph() sess.close() bokeh_session.server_session.close("test complete") os.killpg(os.getpgid(process.pid), signal.SIGTERM)
def test_select_per_layer_comp_ratios(self): pruner = unittest.mock.MagicMock() eval_func = unittest.mock.MagicMock() rounding_algo = unittest.mock.MagicMock() rounding_algo.round.side_effect = [ 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9 ] eval_func.side_effect = [ 10, 20, 30, 40, 50, 60, 70, 80, 90, 11, 21, 31, 35, 40, 45, 50, 55, 60 ] # create tf.compat.v1.Session and initialize the weights and biases with zeros config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True # create session with graph sess = tf.compat.v1.Session(graph=tf.Graph(), config=config) with sess.graph.as_default(): # by default, model will be constructed in default graph _ = mnist_tf_model.create_model(data_format='channels_last') sess.run(tf.compat.v1.global_variables_initializer()) # Create a layer database layer_db = LayerDatabase(model=sess, input_shape=(1, 28, 28, 1), working_dir=None) layer1 = layer_db.find_layer_by_name('conv2d/Conv2D') layer2 = layer_db.find_layer_by_name('conv2d_1/Conv2D') selected_layers = [layer1, layer2] layer_db.mark_picked_layers([layer1, layer2]) try: os.remove('./data/greedy_selection_eval_scores_dict.pkl') except OSError: pass url, process = start_bokeh_server_session(8006) bokeh_session = BokehServerSession(url=url, session_id="compression") # Instantiate child greedy_algo = comp_ratio_select.GreedyCompRatioSelectAlgo( layer_db=layer_db, pruner=pruner, cost_calculator=SpatialSvdCostCalculator(), eval_func=eval_func, eval_iterations=20, cost_metric=CostMetric.mac, target_comp_ratio=Decimal(0.6), num_candidates=10, use_monotonic_fit=True, saved_eval_scores_dict=None, comp_ratio_rounding_algo=rounding_algo, use_cuda=False, bokeh_session=bokeh_session) layer_comp_ratio_list, stats = greedy_algo.select_per_layer_comp_ratios( ) original_cost = SpatialSvdCostCalculator.compute_model_cost(layer_db) for layer in layer_db: if layer not in selected_layers: layer_comp_ratio_list.append(LayerCompRatioPair(layer, None)) compressed_cost = SpatialSvdCostCalculator.calculate_compressed_cost( layer_db, layer_comp_ratio_list, CostMetric.mac) rounding_algo.round.side_effect = [ 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9 ] actual_compression_ratio = compressed_cost.mac / original_cost.mac self.assertTrue( math.isclose(Decimal(0.6), actual_compression_ratio, abs_tol=0.05)) self.assertTrue( os.path.isfile('./data/greedy_selection_eval_scores_dict.pkl')) print('\n') for pair in layer_comp_ratio_list: print(pair) # lets repeat with a saved eval_dict greedy_algo = comp_ratio_select.GreedyCompRatioSelectAlgo( layer_db=layer_db, pruner=pruner, cost_calculator=SpatialSvdCostCalculator(), eval_func=eval_func, eval_iterations=20, cost_metric=CostMetric.mac, target_comp_ratio=Decimal(0.6), num_candidates=10, use_monotonic_fit=True, saved_eval_scores_dict= './data/greedy_selection_eval_scores_dict.pkl', comp_ratio_rounding_algo=rounding_algo, use_cuda=False, bokeh_session=bokeh_session) layer_comp_ratio_list, stats = greedy_algo.select_per_layer_comp_ratios( ) original_cost = SpatialSvdCostCalculator.compute_model_cost(layer_db) for layer in layer_db: if layer not in selected_layers: layer_comp_ratio_list.append(LayerCompRatioPair(layer, None)) compressed_cost = SpatialSvdCostCalculator.calculate_compressed_cost( layer_db, layer_comp_ratio_list, CostMetric.mac) actual_compression_ratio = compressed_cost.mac / original_cost.mac self.assertTrue( math.isclose(Decimal(0.6), actual_compression_ratio, abs_tol=0.05)) print('\n') for pair in layer_comp_ratio_list: print(pair) tf.compat.v1.reset_default_graph() sess.close() bokeh_session.server_session.close("test complete") os.killpg(os.getpgid(process.pid), signal.SIGTERM)
def test_eval_scores_with_spatial_svd_pruner(self): pruner = SpatialSvdPruner() eval_func = unittest.mock.MagicMock() eval_func.side_effect = [ 90, 80, 70, 60, 50, 40, 30, 20, 10, 91, 81, 71, 61, 51, 41, 31, 21, 11 ] # create tf.compat.v1.Session and initialize the weights and biases with zeros config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True # create session with graph sess = tf.compat.v1.Session(graph=tf.Graph(), config=config) with sess.graph.as_default(): # by default, model will be constructed in default graph _ = mnist_tf_model.create_model(data_format='channels_last') sess.run(tf.compat.v1.global_variables_initializer()) # Create a layer database layer_db = LayerDatabase(model=sess, input_shape=(1, 28, 28, 1), working_dir=None) layer1 = layer_db.find_layer_by_name('conv2d/Conv2D') layer2 = layer_db.find_layer_by_name('conv2d_1/Conv2D') layer_db.mark_picked_layers([layer1, layer2]) url, process = start_bokeh_server_session(8006) bokeh_session = BokehServerSession(url=url, session_id="compression") # Instantiate child greedy_algo = comp_ratio_select.GreedyCompRatioSelectAlgo( layer_db=layer_db, pruner=pruner, cost_calculator=SpatialSvdCostCalculator(), eval_func=eval_func, eval_iterations=20, cost_metric=CostMetric.mac, target_comp_ratio=0.5, num_candidates=10, use_monotonic_fit=True, saved_eval_scores_dict=None, comp_ratio_rounding_algo=None, use_cuda=False, bokeh_session=bokeh_session) dict = greedy_algo._compute_eval_scores_for_all_comp_ratio_candidates() print() print(dict) self.assertEqual(90, dict['conv2d/Conv2D'][Decimal('0.1')]) self.assertEqual(51, dict['conv2d_1/Conv2D'][Decimal('0.5')]) self.assertEqual(21, dict['conv2d_1/Conv2D'][Decimal('0.8')]) tf.compat.v1.reset_default_graph() sess.close() bokeh_session.server_session.close("test complete") os.killpg(os.getpgid(process.pid), signal.SIGTERM)
def prune_model(self, layer_db: LayerDatabase, layer_comp_ratio_list: List[LayerCompRatioPair], cost_metric: CostMetric, trainer): # sort all the layers in layer_comp_ratio_list based on occurrence layer_comp_ratio_list = self._sort_on_occurrence( layer_db.model, layer_comp_ratio_list) # Copy the db comp_layer_db = copy.deepcopy(layer_db) current_sess = comp_layer_db.model # Dictionary to map original layer name to list of most recent pruned layer name and output mask. # Masks remain at the original length and specify channels winnowed after each round of winnower. orig_layer_name_to_pruned_name_and_mask_dict = {} # Dictionary to map most recent pruned layer name to the original layer name pruned_name_to_orig_name_dict = {} # List to hold original layers to reconstruct layers_to_reconstruct = [] detached_op_names = set() # Prune layers which have comp ratios less than 1 for layer_comp_ratio in layer_comp_ratio_list: orig_layer = layer_db.find_layer_by_name( layer_comp_ratio.layer.name) if layer_comp_ratio.comp_ratio is not None and layer_comp_ratio.comp_ratio < 1.0: # 1) channel selection prune_indices = self._select_inp_channels( orig_layer, layer_comp_ratio.comp_ratio) if not prune_indices: continue # 2) Winnowing the model current_sess, ordered_modules_list = winnow.winnow_tf_model( current_sess, self._input_op_names, self._output_op_names, [(orig_layer.module, prune_indices)], reshape=self._allow_custom_downsample_ops, in_place=True, verbose=False) if not ordered_modules_list: continue layers_to_reconstruct.append(orig_layer) # Update dictionaries with new info about pruned ops and new masks self._update_pruned_ops_and_masks_info( ordered_modules_list, orig_layer_name_to_pruned_name_and_mask_dict, pruned_name_to_orig_name_dict, detached_op_names) # Save and reload modified graph to allow changes to take effect # Need to initialize uninitialized variables first since only newly winnowed conv ops are initialized during # winnow_tf_model, and all other newly winnowed ops are not. with current_sess.graph.as_default(): initialize_uninitialized_vars(current_sess) current_sess = save_and_load_graph('./saver', current_sess) comp_layer_db.update_database(current_sess, detached_op_names, update_model=True) # Perform reconstruction self._reconstruct_layers(layers_to_reconstruct, orig_layer_name_to_pruned_name_and_mask_dict, layer_db, comp_layer_db) return comp_layer_db
def calculate_compressed_cost( self, layer_db: LayerDatabase, layer_comp_ratio_list: List[LayerCompRatioPair]) -> Cost: """ Calculate cost of a compressed model given a set of layers and corresponding comp-ratios :param layer_db: Layer database for original model :param layer_comp_ratio_list: List of (layer + comp-ratio) pairs :return: Estimated cost of the compressed model """ # sort all the layers in layer_comp_ratio_list based on occurrence layer_comp_ratio_list = self._sort_on_occurrence( layer_db.model, layer_comp_ratio_list) detached_op_names = set() # Copy the db comp_layer_db = copy.deepcopy(layer_db) current_sess = comp_layer_db.model for layer_comp_ratio in layer_comp_ratio_list: orig_layer = layer_db.find_layer_by_name( layer_comp_ratio.layer.name) comp_ratio = layer_comp_ratio.comp_ratio if comp_ratio is not None and comp_ratio < 1.0: # select input channels of conv2d op to winnow prune_indices = self._select_inp_channels( orig_layer, comp_ratio) if not prune_indices: continue # Winnow the selected op and modify it's upstream affected ops current_sess, ordered_modules_list = winnow.winnow_tf_model( current_sess, self._input_op_names, self._output_op_names, [(orig_layer.module, prune_indices)], reshape=self._allow_custom_downsample_ops, in_place=True, verbose=False) if not ordered_modules_list: continue # Get all the detached op names from updated session graph for orig_op_name, _, _, _ in ordered_modules_list: detached_op_names.add(orig_op_name) # update layer database by excluding the detached ops comp_layer_db.update_database(current_sess, detached_op_names, update_model=False) # calculate the cost of this model compressed_model_cost = CostCalculator.compute_model_cost( comp_layer_db) # close the session associated with compressed layer database comp_layer_db.model.close() return compressed_model_cost
def test_datasampling_and_reconstruction(self): """ Test data sampling and reconstruction logic """ tf.compat.v1.reset_default_graph() batch_size = 1 input_data = np.random.rand(100, 224, 224, 3) dataset = tf.data.Dataset.from_tensor_slices(input_data) dataset = dataset.batch(batch_size=batch_size) orig_g = tf.Graph() with orig_g.as_default(): _ = VGG16(weights=None, input_shape=(224, 224, 3), include_top=False) orig_init = tf.compat.v1.global_variables_initializer() input_op_names = ['input_1'] output_op_names = ['block5_pool/MaxPool'] # create sess with graph orig_sess = tf.compat.v1.Session(graph=orig_g) # initialize all the variables in VGG16 orig_sess.run(orig_init) # create layer database layer_db = LayerDatabase(model=orig_sess, input_shape=(1, 224, 224, 3), working_dir=None) conv_layer = layer_db.find_layer_by_name('block1_conv1/Conv2D') comp_layer_db = copy.deepcopy(layer_db) comp_conv_layer = comp_layer_db.find_layer_by_name( 'block1_conv1/Conv2D') # get the weights before reconstruction in original model before_recon_weights_orig_model = layer_db.model.run( conv_layer.module.inputs[1]) # get the weights before reconstruction in pruned model before_recon_weights_pruned_model = comp_layer_db.model.run( comp_conv_layer.module.inputs[1]) # weight should be exactly same before reconstruction in original and pruned layer database self.assertTrue( np.array_equal(before_recon_weights_orig_model, before_recon_weights_pruned_model)) cp = InputChannelPruner(input_op_names=input_op_names, output_op_names=output_op_names, data_set=dataset, batch_size=batch_size, num_reconstruction_samples=50, allow_custom_downsample_ops=True) num_in_channels = comp_conv_layer.weight_shape[0] cp._data_subsample_and_reconstruction(orig_layer=conv_layer, pruned_layer=comp_conv_layer, output_mask=[1] * num_in_channels, orig_layer_db=layer_db, comp_layer_db=comp_layer_db) # get the weights after reconstruction after_recon_weights_pruned_model = comp_layer_db.model.run( comp_conv_layer.module.inputs[1]) # weight should not be same before and after reconstruction self.assertFalse( np.array_equal(before_recon_weights_orig_model, after_recon_weights_pruned_model)) layer_db.model.close() comp_layer_db.model.close() # delete temp directory shutil.rmtree(str('./temp_meta/'))